import gradio as gr
from convert_url_to_diffusers_sdxl_gr import (
    convert_url_to_diffusers_repo,
    SCHEDULER_CONFIG_MAP,
)

vaes = [
    "",
    "https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/blob/main/sdxl.vae.safetensors",
    "https://huggingface.co/nubby/blessed-sdxl-vae-fp16-fix/blob/main/sdxl_vae-fp16fix-blessed.safetensors",
]
loras = [
    "",
    "https://huggingface.co/SPO-Diffusion-Models/SPO-SDXL_4k-p_10ep_LoRA/blob/main/spo_sdxl_10ep_4k-data_lora_diffusers.safetensors",
    "https://huggingface.co/wangfuyun/PCM_Weights/blob/main/sdxl/pcm_sdxl_smallcfg_2step_converted.safetensors",
    "https://huggingface.co/wangfuyun/PCM_Weights/blob/main/sdxl/pcm_sdxl_smallcfg_4step_converted.safetensors",
    "https://huggingface.co/wangfuyun/PCM_Weights/blob/main/sdxl/pcm_sdxl_smallcfg_8step_converted.safetensors",
    "https://huggingface.co/wangfuyun/PCM_Weights/blob/main/sdxl/pcm_sdxl_normalcfg_8step_converted.safetensors",
    "https://huggingface.co/wangfuyun/PCM_Weights/blob/main/sdxl/pcm_sdxl_normalcfg_16step_converted.safetensors",
    "https://huggingface.co/ByteDance/Hyper-SD/blob/main/Hyper-SDXL-1step-lora.safetensors",
    "https://huggingface.co/ByteDance/Hyper-SD/blob/main/Hyper-SDXL-2steps-lora.safetensors",
    "https://huggingface.co/ByteDance/Hyper-SD/blob/main/Hyper-SDXL-4steps-lora.safetensors",
    "https://huggingface.co/ByteDance/Hyper-SD/blob/main/Hyper-SDXL-8steps-CFG-lora.safetensors",
    "https://huggingface.co/ByteDance/Hyper-SD/blob/main/Hyper-SDXL-12steps-CFG-lora.safetensors",
    "https://huggingface.co/latent-consistency/lcm-lora-sdxl/blob/main/pytorch_lora_weights.safetensors",
]
schedulers = list(SCHEDULER_CONFIG_MAP.keys())

css = """"""

with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as demo:
    gr.Markdown("# Download and convert any Stable Diffusion XL safetensors to Diffusers and create your repo")
    gr.Markdown(
        f"""
**⚠️IMPORTANT NOTICE⚠️**<br>
From an information security standpoint, it is dangerous to expose your access token or key to others.
If you do use it, I recommend that you duplicate this space on your own account before doing so.
Keys and tokens could be set to SECRET (HF_TOKEN, CIVITAI_API_KEY) if it's placed in your own space.
It saves you the trouble of typing them in.<br>
<br>
**The steps are the following**:
- Paste a write-access token from [hf.co/settings/tokens](https://huggingface.co/settings/tokens).
- Input a model download url from the Hub or Civitai or other sites.
- If you want to download a model from Civitai, paste a Civitai API Key.
- Input your new repo name. e.g. 'yourid/newrepo'.
- Set the parameters. If not sure, just use the defaults.
- Click "Submit".
- Patiently wait until the output changes.
            """
    )
    with gr.Column():
        dl_url = gr.Textbox(label="URL to download", placeholder="https://...", value="", max_lines=1)
        repo_id = gr.Textbox(label="Your New Repo ID", placeholder="author/model", value="", max_lines=1)
        hf_token = gr.Textbox(label="Your HF write token", placeholder="", value="", max_lines=1)
        civitai_key = gr.Textbox(label="Your Civitai API Key (Optional)", info="If you download model from Civitai...", placeholder="", value="", max_lines=1)
        is_half = gr.Checkbox(label="Half precision", value=True)
        vae = gr.Dropdown(label="VAE", choices=vaes, value="", allow_custom_value=True)
        scheduler = gr.Dropdown(label="Scheduler (Sampler)", choices=schedulers, value="Euler a")
        lora1 = gr.Dropdown(label="LoRA1", choices=loras, value="", allow_custom_value=True)
        lora1s = gr.Slider(minimum=-2, maximum=2, step=0.01, value=1.00, label="LoRA1 weight scale")
        lora2 = gr.Dropdown(label="LoRA2", choices=loras, value="", allow_custom_value=True)
        lora2s = gr.Slider(minimum=-2, maximum=2, step=0.01, value=1.00, label="LoRA2 weight scale")
        lora3 = gr.Dropdown(label="LoRA3", choices=loras, value="", allow_custom_value=True)
        lora3s = gr.Slider(minimum=-2, maximum=2, step=0.01, value=1.00, label="LoRA3 weight scale")
        lora4 = gr.Dropdown(label="LoRA4", choices=loras, value="", allow_custom_value=True)
        lora4s = gr.Slider(minimum=-2, maximum=2, step=0.01, value=1.00, label="LoRA4 weight scale")
        lora5 = gr.Dropdown(label="LoRA5", choices=loras, value="", allow_custom_value=True)
        lora5s = gr.Slider(minimum=-2, maximum=2, step=0.01, value=1.00, label="LoRA5 weight scale")
        run_button = gr.Button(value="Submit")
        repo_urls = gr.CheckboxGroup(visible=False, choices=[], value=None)
        output_md = gr.Markdown(label="Output")

    gr.on(
        triggers=[run_button.click],
        fn=convert_url_to_diffusers_repo,
        inputs=[dl_url, repo_id, hf_token, civitai_key, repo_urls, is_half, vae, scheduler,
                 lora1, lora1s, lora2, lora2s, lora3, lora3s, lora4, lora4s, lora5, lora5s],
        outputs=[repo_urls, output_md],
    )

demo.queue()
demo.launch()