up

app.py

@@ -1,198 +1,36 @@
-import os
-import subprocess
-from huggingface_hub import HfApi, upload_folder
 import gradio as gr
-import hf_utils
-import utils
-from safetensors import safe_open
-import torch
 
-subprocess.run(["git", "clone", "https://github.com/huggingface/diffusers", "diffs"])
-
-def error_str(error, title="Error"):
-    return f"""#### {title}
-            {error}"""  if error else ""
-
-def on_token_change(token):
-    model_names, error = hf_utils.get_my_model_names(token)
-    if model_names:
-        model_names.append("Other")
-
-    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))
-
-def url_to_model_id(model_id_str):
-    return model_id_str.split("/")[-2] + "/" + model_id_str.split("/")[-1] if model_id_str.startswith("https://huggingface.co/") else model_id_str
-    
-def get_ckpt_names(token, radio_model_names, input_model):
-    
-    model_id = url_to_model_id(input_model) if radio_model_names == "Other" else radio_model_names
-
-    if token == "" or model_id == "":
-        return error_str("Please enter both a token and a model name.", title="Invalid input"), gr.update(choices=[]), gr.update(visible=False)
-
-    try:
-        api = HfApi(token=token)
-        ckpt_files = [f for f in api.list_repo_files(repo_id=model_id) if f.endswith(".ckpt") or f.endswith(".safetensors")]
-        
-        if not ckpt_files:
-            return error_str("No checkpoint files found in the model repo."), gr.update(choices=[]), gr.update(visible=False)
-        
-        return None, gr.update(choices=ckpt_files, value=ckpt_files[0], visible=True), gr.update(visible=True)
-        
-    except Exception as e:
-        return error_str(e), gr.update(choices=[]), None
-
-def convert_and_push(radio_model_names, input_model, ckpt_name, sd_version, token, path_in_repo, ema, safetensors):
-    extract_ema = ema == "ema"
-    
-    if sd_version == None:
-        return error_str("You must select a stable diffusion version.", title="Invalid input")
-
-    model_id = url_to_model_id(input_model) if radio_model_names == "Other" else radio_model_names
-
-    try:
-        model_id = url_to_model_id(model_id)
-
-        # 1. Download the checkpoint file
-        ckpt_path, revision = hf_utils.download_file(repo_id=model_id, filename=ckpt_name, token=token)
-
-        if safetensors == "yes":
-            tensors = {}
-            with safe_open(ckpt_path, framework="pt", device="cpu") as f:
-               for key in f.keys():
-                   tensors[key] = f.get_tensor(key)
-
-            new_checkpoint_path = "/".join(ckpt_path.split("/")[:-1] + ["model_safe.ckpt"])
-            torch.save(tensors, new_checkpoint_path)
-            ckpt_path = new_checkpoint_path
-            print("Converting ckpt_path", ckpt_path)
-
-        print(ckpt_path)
-
-        # 2. Run the conversion script
-        os.makedirs(model_id, exist_ok=True)
-        run_command = [
-            "python3",
-            "./diffs/scripts/convert_original_stable_diffusion_to_diffusers.py",
-            "--checkpoint_path",
-            ckpt_path,
-            "--dump_path" ,
-            model_id,
-        ]
-        if extract_ema:
-            run_command.append("--extract_ema")
-        subprocess.run(run_command)
-
-        # 3. Push to the model repo
-        commit_message="Add Diffusers weights"
-        upload_folder(
-            folder_path=model_id,
-            repo_id=model_id,
-            path_in_repo=path_in_repo,
-            token=token,
-            create_pr=True,
-            commit_message=commit_message,
-            commit_description=f"Add Diffusers weights converted from checkpoint `{ckpt_name}` in revision {revision}",
-        )
-
-        # # 4. Delete the downloaded checkpoint file, yaml files, and the converted model folder
-        hf_utils.delete_file(revision)
-        subprocess.run(["rm", "-rf", model_id.split('/')[0]])
-        import glob
-        for f in glob.glob("*.yaml*"):
-            subprocess.run(["rm", "-rf", f])
-
-        return f"""Successfully converted the checkpoint and opened a PR to add the weights to the model repo.
-                You can view and merge the PR [here]({hf_utils.get_pr_url(HfApi(token=token), model_id, commit_message)})."""
-
-        return "Done"
-    
-    except Exception as e:
-        return error_str(e)
-
-
-DESCRIPTION = """### Convert a stable diffusion checkpoint to Diffusers🧨
-                With this space, you can easily convert a CompVis stable diffusion checkpoint to Diffusers and automatically create a pull request to the model repo.
-                You can choose to convert a checkpoint from one of your own models, or from any other model on the Hub.
-                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)"""
-
-with gr.Blocks() as demo:
-
-    gr.Markdown(DESCRIPTION)
-    with gr.Row():
-
-        with gr.Column(scale=11):
-            with gr.Column():
-                gr.Markdown("## 1. Load model info")
-                input_token = gr.Textbox(
-                    max_lines=1,
-                    type="password",
-                    label="Enter your Hugging Face token",
-                    placeholder="READ permission is sufficient"
-                )
-                gr.Markdown("You can get a token [here](https://huggingface.co/settings/tokens)")
-                with gr.Group(visible=False) as group_model:
-                    radio_model_names = gr.Radio(label="Choose a model")
-                    input_model = gr.Textbox(
-                        max_lines=1,
-                        label="Model name or URL",
-                        placeholder="username/model_name",
-                        visible=False,
-                    )
-
-            btn_get_ckpts = gr.Button("Load", visible=False)
-
-        with gr.Column(scale=10):
-            with gr.Column(visible=False) as group_convert:
-                gr.Markdown("## 2. Convert to Diffusers🧨")
-                radio_ckpts = gr.Radio(label="Choose the checkpoint to convert", visible=False)
-                path_in_repo = gr.Textbox(label="Path where the weights will be saved", placeholder="Leave empty for root folder")
-                ema = gr.Radio(label="Extract EMA or non-EMA?", choices=["ema", "non-ema"])
-                safetensors = gr.Radio(label="Extract from safetensors", choices=["yes", "no"], value="no")
-                radio_sd_version = gr.Radio(label="Choose the model version", choices=["v1", "v2", "v2.1"])
-                gr.Markdown("Conversion may take a few minutes.")
-                btn_convert = gr.Button("Convert & Push")
-
-    error_output = gr.Markdown(label="Output")
-
-    input_token.change(
-        fn=on_token_change,
-        inputs=input_token,
-        outputs=[group_model, radio_model_names, btn_get_ckpts, error_output],
-        queue=False,
-        scroll_to_output=True)
-    
-    radio_model_names.change(
-        lambda x: gr.update(visible=x == "Other"),
-        inputs=radio_model_names,
-        outputs=input_model,
-        queue=False,
-        scroll_to_output=True)
-    
-    btn_get_ckpts.click(
-        fn=get_ckpt_names,
-        inputs=[input_token, radio_model_names, input_model],
-        outputs=[error_output, radio_ckpts, group_convert],
-        scroll_to_output=True,
-        queue=False
-    )
-
-    btn_convert.click(
-        fn=convert_and_push,
-        inputs=[radio_model_names, input_model, radio_ckpts, radio_sd_version, input_token, path_in_repo, ema, safetensors],
-        outputs=error_output,
-        scroll_to_output=True
-    )
-
-    # gr.Markdown("""<img src="https://raw.githubusercontent.com/huggingface/diffusers/main/docs/source/imgs/diffusers_library.jpg" width="150"/>""")
-    gr.HTML("""
-    <div style="border-top: 1px solid #303030;">
-      <br>
-      <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>
-      <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>
-      <p><img src="https://visitor-badge.glitch.me/badge?page_id=anzorq.sd-to-diffusers" alt="visitors"></p>
-    </div>
-    """)
-    
-demo.queue()
-demo.launch(debug=True, share=utils.is_google_colab())
+from convert import convert
+
+DESCRIPTION = """
+The steps are the following:
+
+- Paste a read-access token from hf.co/settings/tokens. Read access is enough given that we will open a PR against the source repo.
+- Input a model id from the Hub
+- Input the filename from the root dir of the repo that you would like to convert, e.g. 'v2-1_768-ema-pruned.ckpt' or 'v1-5-pruned.safetensors'
+- Chose which Stable Diffusion version, image size, scheduler type the model has and whether you want the "ema", or "non-ema" weights.
+- Click "Submit"
+- That's it! You'll get feedback if it works or not, and if it worked, you'll get the URL of the opened PR 🔥
+
+⚠️ For now only `pytorch_model.bin` files are supported but we'll extend in the future.
+"""
+
+demo = gr.Interface(
+    title="Convert any model to Safetensors and open a PR",
+    description=DESCRIPTION,
+    allow_flagging="never",
+    article="Check out the [Safetensors repo on GitHub](https://github.com/huggingface/safetensors)",
+    inputs=[
+        gr.Text(max_lines=1, label="your_hf_token"),
+        gr.Text(max_lines=1, label="model_id"),
+        gr.Text(max_lines=1, label="filename"),
+        gr.Radio(label="Model type", choices=["v1", "v2.0", "v2.1"]),
+        gr.Radio(label="Sample size (px)", choices=[512, 768]),
+        gr.Radio(label="Scheduler type", choices=["pndm", "heun", "euler", "dpm", "ddim"], value="dpm"),
+        gr.Radio(label="Extract EMA or non-EMA?", choices=["ema", "non-ema"], value="ema"),
+    ],
+    outputs=[gr.Markdown(label="output")],
+    fn=convert,
+).queue(max_size=10, concurrency_count=1)
+
+demo.launch(show_api=True)

convert.py

@@ -0,0 +1,79 @@
+import argparse
+import requests
+import json
+import os
+import shutil
+from collections import defaultdict
+from inspect import signature
+from tempfile import TemporaryDirectory
+from typing import Dict, List, Optional, Set
+
+import torch
+from io import BytesIO
+
+from huggingface_hub import CommitInfo, Discussion, HfApi, hf_hub_download
+from huggingface_hub.file_download import repo_folder_name
+from diffusers.pipelines.stable_diffusion.convert_from_ckpt import download_from_original_stable_diffusion_ckpt
+from transformers import CONFIG_MAPPING
+
+
+COMMIT_MESSAGE = " This PR adds the both fp32 and fp16 in PyTorch and safetensors format to {}"
+
+
+def convert_single(model_id: str, filename: str, model_type: str, sample_size: int, scheduler_type: str, extract_ema: bool, folder: str):
+    from_safetensors = filename.endswith(".safetensors")
+
+    local_file = os.path.join(model_id, filename)
+    ckpt_file = local_file if os.path.isfile(local_file) else hf_hub_download(repo_id=model_id, filename=filename)
+
+    if model_type == "v1":
+        config_url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml"
+    elif model_type == "v2.0":
+        config_url = "https://raw.githubusercontent.com/Stability-AI/stablediffusion/main/configs/stable-diffusion/v2-inference.yaml"
+    elif model_type == "v2.1":
+        config_url = "https://raw.githubusercontent.com/Stability-AI/stablediffusion/main/configs/stable-diffusion/v2-inference-v.yaml"
+
+    config_file = BytesIO(requests.get(config_url).content)
+
+    pipeline = download_from_original_stable_diffusion_ckpt(ckpt_file, config_file, image_size=sample_size, scheduler_type=scheduler_type, from_safetensors=from_safetensors, extract_ema=extract_ema)
+
+    pipeline.save_pretrained(folder)
+    pipeline.save_pretrained(folder, safe_serialization=True)
+
+    pipeline = pipeline.to(torch_dtype=torch.float16)
+    pipeline.save_pretrained(folder, variant="fp16")
+    pipeline.save_pretrained(folder, safe_serialization=True, variant="fp16")
+
+    return folder
+
+
+def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discussion"]:
+    try:
+        discussions = api.get_repo_discussions(repo_id=model_id)
+    except Exception:
+        return None
+    for discussion in discussions:
+        if discussion.status == "open" and discussion.is_pull_request and discussion.title == pr_title:
+            details = api.get_discussion_details(repo_id=model_id, discussion_num=discussion.num)
+            if details.target_branch == "refs/heads/main":
+                return discussion
+
+
+def convert(token: str, model_id: str, filename: str, model_type: str, sample_size: int = 512, scheduler_type: str = "pndm", extract_ema: bool = True):
+    api = HfApi()
+
+    pr_title = "Adding `diffusers` weights of this model"
+
+    with TemporaryDirectory() as d:
+        folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models"))
+        os.makedirs(folder)
+        new_pr = None
+        try:
+            folder = convert_single(model_id, filename, model_type, sample_size, scheduler_type, extract_ema, folder)
+            new_pr  = api.upload_folder(folder_path=folder, path_in_repo="./", repo_id=model_id, repo_type="model", token=token, commit_description=COMMIT_MESSAGE.format(model_id), create_pr=True)
+            pr_number = new_pr.split("%2F")[-1].split("/")[0]
+            print(f"Pr created at: {'https://huggingface.co/' + os.path.join(model_id, 'discussions', pr_number)}")
+        finally:
+            shutil.rmtree(folder)
+
+        return new_pr

requirements.txt

@@ -1,10 +1,7 @@
-git+https://github.com/huggingface/huggingface_hub@main
-git+https://github.com/huggingface/diffusers.git
---find-links https://download.pytorch.org/whl/cpu
-torch
-#transformers
-git+https://github.com/huggingface/transformers
-lightning
-OmegaConf
-ftfy
-safetensors
+huggingface_hub
+safetensors
+transformers
+accelerate
+git+https://github.com/huggingface/diffusers
+omegaconf
+pytorch_lightning