import gradio as gr
from merge_gr import gen_repo_list, upload_repo_list, clear_repo_list, process_repos_gr

css = """"""

with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as demo:
    gr.Markdown("# SDXL/SD1.5 DARE Merger (experiment)")
    gr.Markdown(
        f"""
This Space is a mod version of [martyn](https://huggingface.co/martyn)'s [safetensors-merge-supermario](https://github.com/martyn/safetensors-merge-supermario) forced to be compatible with Diffusers.
Since the space is completely experimental and unfinished,
I recommend using [ComfyUI-DareMerge](https://github.com/54rt1n/ComfyUI-DareMerge)
 or [WebUI SuperMerger](https://github.com/hako-mikan/sd-webui-supermerger) for actual merging.
 Also, I think most safetensors models with the same structure can be merged even if they are not SD models, but I haven't tried.<br>
**⚠️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 running.
Keys and tokens could be set to SECRET (HF_TOKEN) 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.
- Input your HF user ID. e.g. 'yourid'.
- Input your new merged repo name.
- Input information for merging models.
- Set the parameters. If not sure, just use the defaults.
- Click "Submit".
- Patiently wait until the output changes. It takes approximately 5~6 minutes (downloading from HF).
            """
    )
    with gr.Column():
        with gr.Group():
            hf_token = gr.Textbox(label="Your HF write token", placeholder="hf_...", value="", max_lines=1)
            with gr.Row():
                hf_user = gr.Textbox(label="Your HF user ID", placeholder="username", value="", max_lines=1)
                hf_repo = gr.Textbox(label="New repo name", placeholder="reponame", value="", max_lines=1)
        with gr.Group():
            with gr.Accordion("YAML", open=True):
                merge_yaml_input = gr.Textbox(label="List of Repos or URLs to merge", placeholder="author/repo\nauthor/repo\n...", value="", lines=4)
                merge_yaml_md = gr.Markdown()
                merge_yaml_upload = gr.UploadButton(label="Upload YAML file", file_types=["text"])
                merge_yaml_clear = gr.Button("Clear YAML files")
            with gr.Row():
                merge_p = gr.Number(label="Default dropout probability", value=0.5, minimum=0, maximum=1.0, step=0.01)
                merge_lambda = gr.Number(label="Default scaling factor for the weight delta", value=1.0, minimum=0, maximum=2.0, step=0.01)
            merge_mode = gr.Radio(label="Mode", choices=["SDXL", "SD1.5", "Single files"], value="SDXL")
            merge_is_upload_sf = gr.Checkbox(label="Convert Diffusers files to single safetensors file", value=False)
            merge_skip = gr.CheckboxGroup(label="Skip Diffusers folders", choices=["vae", "text_encoder", "text_encoder_2", "text_encoder_3"], value=["vae", "text_encoder"])
            merge_is_upload = gr.Checkbox(label="Upload files into new Repo", value=True)
            merge_repo_exists_ok = gr.Checkbox(label="Overwrite exsisting Repo", value=False)
        run_button = gr.Button(value="Submit")
        repo_urls = gr.CheckboxGroup(visible=False, choices=[], value=None)
        output_md = gr.Markdown(label="Output")
        merge_files = gr.Files(label="Download", interactive=False, value=[])

    merge_yaml_input.change(gen_repo_list, [merge_yaml_input, merge_p, merge_lambda], [merge_yaml_md], queue=False)
    merge_yaml_upload.upload(upload_repo_list, [merge_yaml_upload, merge_p, merge_lambda], [merge_yaml_md], queue=False)
    merge_yaml_clear.click(clear_repo_list, None, [merge_yaml_input, merge_yaml_md], queue=False)
    gr.on(
        triggers=[run_button.click],
        fn=process_repos_gr,
        inputs=[merge_mode, merge_p, merge_lambda, merge_skip, hf_user, hf_repo, hf_token,
                 merge_is_upload, merge_is_upload_sf, merge_repo_exists_ok, merge_files, repo_urls],
        outputs=[merge_files, repo_urls, output_md],
    )

demo.queue()
demo.launch()