sophesrex/sd-to-diffusers

1import os2import subprocess3from huggingface_hub import HfApi, upload_folder4import gradio as gr5import hf_utils6import utils7 8subprocess.run(["git", "clone", "https://github.com/huggingface/diffusers.git", "diffs"])9 10def error_str(error, title="Error"):11    return f"""#### {title}12            {error}"""  if error else ""13 14def on_token_change(token):15    model_names, error = hf_utils.get_my_model_names(token)16    if model_names:17        model_names.append("Other")18 19    return gr.update(visible=bool(model_names)), gr.update(choices=model_names, value=model_names[0] if model_names else None), gr.update(visible=bool(model_names)), gr.update(value=error_str(error))20 21def url_to_model_id(model_id_str):22    return model_id_str.split("/")[-2] + "/" + model_id_str.split("/")[-1] if model_id_str.startswith("https://huggingface.co/") else model_id_str23    24def get_ckpt_names(token, radio_model_names, input_model):25    26    model_id = url_to_model_id(input_model) if radio_model_names == "Other" else radio_model_names27 28    if token == "" or model_id == "":29        return error_str("Please enter both a token and a model name.", title="Invalid input"), gr.update(choices=[]), gr.update(visible=False)30 31    try:32        api = HfApi(token=token)33        ckpt_files = [f for f in api.list_repo_files(repo_id=model_id) if f.endswith(".ckpt")]34        35        if not ckpt_files:36            return error_str("No checkpoint files found in the model repo."), gr.update(choices=[]), gr.update(visible=False)37        38        return None, gr.update(choices=ckpt_files, value=ckpt_files[0], visible=True), gr.update(visible=True)39        40    except Exception as e:41        return error_str(e), gr.update(choices=[]), None42 43def convert_and_push(radio_model_names, input_model, ckpt_name, token, path_in_repo):44    45    model_id = url_to_model_id(input_model) if radio_model_names == "Other" else radio_model_names46 47    try:48        model_id = url_to_model_id(model_id)49 50        # 1. Download the checkpoint file51        ckpt_path, revision = hf_utils.download_file(repo_id=model_id, filename=ckpt_name, token=token)52 53        # 2. Run the conversion script54        os.makedirs(model_id, exist_ok=True)55        subprocess.run(56            [57                "python3",58                "./diffs/scripts/convert_original_stable_diffusion_to_diffusers.py",59                "--checkpoint_path",60                ckpt_path,61                "--dump_path" ,62                model_id,63            ]64        )65 66        # 3. Push to the model repo67        commit_message="Add Diffusers weights"68        upload_folder(69            folder_path=model_id,70            repo_id=model_id,71            path_in_repo=path_in_repo,72            token=token,73            create_pr=True,74            commit_message=commit_message,75            commit_description=f"Add Diffusers weights converted from checkpoint `{ckpt_name}` in revision {revision}",76        )77 78        # # 4. Delete the downloaded checkpoint file, yaml files, and the converted model folder79        hf_utils.delete_file(revision)80        subprocess.run(["rm", "-rf", model_id.split('/')[0]])81        import glob82        for f in glob.glob("*.yaml*"):83            subprocess.run(["rm", "-rf", f])84 85        return f"""Successfully converted the checkpoint and opened a PR to add the weights to the model repo.86                You can view and merge the PR [here]({hf_utils.get_pr_url(HfApi(token=token), model_id, commit_message)})."""87    88    except Exception as e:89        return error_str(e)90 91 92DESCRIPTION = """### Convert a stable diffusion checkpoint to Diffusers🧨93                With this space, you can easily convert a CompVis stable diffusion checkpoint to Diffusers and automatically create a pull request to the model repo.94                You can choose to convert a checkpoint from one of your own models, or from any other model on the Hub.95                You can skip the queue by running the app in the colab: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/gist/qunash/f0f3152c5851c0c477b68b7b98d547fe/convert-sd-to-diffusers.ipynb)"""96 97with gr.Blocks() as demo:98 99    gr.Markdown(DESCRIPTION)100    with gr.Row():101 102        with gr.Column(scale=11):103            with gr.Column():104                gr.Markdown("## 1. Load model info")105                input_token = gr.Textbox(106                    max_lines=1,107                    label="Enter your Hugging Face token",108                    placeholder="READ permission is enough",109                )110                gr.Markdown("You can get a token [here](https://huggingface.co/settings/tokens)")111                with gr.Group(visible=False) as group_model:112                    radio_model_names = gr.Radio(label="Choose a model")113                    input_model = gr.Textbox(114                        max_lines=1,115                        label="Model name or URL",116                        placeholder="username/model_name",117                        visible=False,118                    )119 120            btn_get_ckpts = gr.Button("Load", visible=False)121 122        with gr.Column(scale=10):123            with gr.Column(visible=False) as group_convert:124                gr.Markdown("## 2. Convert to Diffusers🧨")125                radio_ckpts = gr.Radio(label="Choose the checkpoint to convert", visible=False)126                path_in_repo = gr.Textbox(label="Path where the weights will be saved", placeholder="Leave empty for root folder")127                gr.Markdown("Conversion may take a few minutes.")128                btn_convert = gr.Button("Convert & Push")129 130    error_output = gr.Markdown(label="Output")131 132    input_token.change(133        fn=on_token_change,134        inputs=input_token,135        outputs=[group_model, radio_model_names, btn_get_ckpts, error_output],136        queue=False,137        scroll_to_output=True)138    139    radio_model_names.change(140        lambda x: gr.update(visible=x == "Other"),141        inputs=radio_model_names,142        outputs=input_model,143        queue=False,144        scroll_to_output=True)145    146    btn_get_ckpts.click(147        fn=get_ckpt_names,148        inputs=[input_token, radio_model_names, input_model],149        outputs=[error_output, radio_ckpts, group_convert],150        scroll_to_output=True,151        queue=False152    )153 154    btn_convert.click(155        fn=convert_and_push,156        inputs=[radio_model_names, input_model, radio_ckpts, input_token, path_in_repo],157        outputs=error_output,158        scroll_to_output=True159    )160 161    # gr.Markdown("""<img src="https://raw.githubusercontent.com/huggingface/diffusers/main/docs/source/imgs/diffusers_library.jpg" width="150"/>""")162    gr.HTML("""163    <div style="border-top: 1px solid #303030;">164      <br>165      <p>Space by: <a href="https://twitter.com/hahahahohohe"><img src="https://img.shields.io/twitter/follow/hahahahohohe?label=%40anzorq&style=social" alt="Twitter Follow"></a></p><br>166      <a href="https://www.buymeacoffee.com/anzorq" target="_blank"><img src="https://cdn.buymeacoffee.com/buttons/v2/default-yellow.png" alt="Buy Me A Coffee" style="height: 45px !important;width: 162px !important;" ></a><br><br>167      <p><img src="https://visitor-badge.glitch.me/badge?page_id=anzorq.sd-to-diffusers" alt="visitors"></p>168    </div>169    """)170    171demo.queue()172demo.launch(share=utils.is_google_colab())173