Add quantize

app.py

@@ -1,7 +1,98 @@
+import csv
+from datetime import datetime
+import os
+from typing import Optional
 import gradio as gr
 
-def greet(name):
-    return "Hello " + name + "!!"
+import quantize
+from huggingface_hub import HfApi, Repository
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+
+DATASET_REPO_URL = "https://huggingface.co/datasets/safetensors/conversions"
+DATA_FILENAME = "data.csv"
+DATA_FILE = os.path.join("data", DATA_FILENAME)
+
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+repo: Optional[Repository] = None
+# TODO
+if False and HF_TOKEN:
+    repo = Repository(local_dir="data", clone_from=DATASET_REPO_URL, token=HF_TOKEN)
+
+
+def run(model_id: str, is_private: bool, token: Optional[str] = None) -> str:
+    if model_id == "":
+        return """
+        ### Invalid input 🐞
+        
+        Please fill a token and model_id.
+        """
+    try:
+        if is_private:
+            api = HfApi(token=token)
+        else:
+            api = HfApi(token=HF_TOKEN)
+        hf_is_private = api.model_info(repo_id=model_id).private
+        if is_private and not hf_is_private:
+            # This model is NOT private
+            # Change the token so we make the PR on behalf of the bot.
+            api = HfApi(token=HF_TOKEN)
+
+        print("is_private", is_private)
+
+        commit_info, errors = quantize.quantize(api=api, model_id=model_id)
+        print("[commit_info]", commit_info)
+
+
+        string =  f"""
+        ### Success 🔥
+        Yay! This model was successfully converted and a PR was open using your token, here:
+        [{commit_info.pr_url}]({commit_info.pr_url})
+        """
+        if errors:
+            string += "\nErrors during conversion:\n"
+            string += "\n".join(f"Error while converting {filename}: {e}, skipped conversion" for filename, e in errors)
+        return string
+    except Exception as e:
+        return f"""
+        ### Error 😢😢😢
+        
+        {e}
+        """
+
+
+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
+- 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.
+"""
+
+title="Quantize model and convert to CoreML"
+allow_flagging="never"
+
+def token_text(visible=False):
+    return gr.Text(max_lines=1, label="your_hf_token", visible=True, value="")
+
+with gr.Blocks(title=title) as demo:
+    description = gr.Markdown(f"""# {title}""")
+    description = gr.Markdown(DESCRIPTION)
+
+    with gr.Row() as r:
+        with gr.Column() as c:
+            model_id = gr.Text(max_lines=1, label="model_id", value="jblalock30/coreml")
+            is_private = gr.Checkbox(label="Private model")
+            token = token_text()
+            with gr.Row() as c:
+                clean = gr.ClearButton()
+                submit = gr.Button("Submit", variant="primary")
+
+        with gr.Column() as d:
+            output = gr.Markdown(value="hi")
+
+    is_private.change(lambda s: token_text(s), inputs=is_private, outputs=token)
+    submit.click(run, inputs=[model_id, is_private, token], outputs=output, concurrency_limit=1)
+
+demo.queue(max_size=10).launch(show_api=True)

ml-stable-diffusion

@@ -0,0 +1 @@
+Subproject commit f7d2f5e9fb0681b15770943e492bf2f6dc3414f3

quantize.py

@@ -0,0 +1,148 @@
+import argparse
+import json
+import os
+import shutil
+from collections import defaultdict
+from tempfile import TemporaryDirectory
+from typing import Dict, List, Optional, Set, Tuple
+import subprocess
+
+import torch
+
+from huggingface_hub import CommitInfo, CommitOperationAdd, Discussion, HfApi, hf_hub_download
+from huggingface_hub.file_download import repo_folder_name
+from safetensors.torch import _find_shared_tensors, _is_complete, load_file, save_file
+
+ConversionResult = Tuple[List["CommitOperationAdd"], List[Tuple[str, "Exception"]]]
+
+def convert_generic(
+    model_id: str, *, revision=Optional[str], folder: str, filenames: Set[str], token: Optional[str]
+) -> ConversionResult:
+    operations = []
+    errors = []
+
+    # python3 -m python_coreml_stable_diffusion.torch2coreml \
+    # --model-version stabilityai/sdxl-turbo \
+    # -o packages/sdxl-turbo \
+    # --convert-unet --convert-text-encoder --convert-vae-decoder --chunk-unet --attention-implementation ORIGINAL \
+    # --bundle-resources-for-swift-cli \
+    # --quantize-nbits 2
+
+ 
+    print("Starting conversion") #
+
+    subprocess.run(["python3", "-m" , "python_coreml_stable_diffusion.torch2coreml", "--model-version", "stabilityai/sd-turbo", "-o", folder, "--convert-unet", "--convert-text-encoder", "--convert-vae-decoder", "--chunk-unet", "--attention-implementation", "ORIGINAL", "--bundle-resources-for-swift-cli"])
+    # with open(f'{folder}/newfile.txt', 'w') as f:
+    #     f.write('Hello, World!')
+
+    print("Done")
+
+    operations.append(CommitOperationAdd(path_in_repo='Resources', path_or_fileobj=f'{folder}/Resources'))
+
+    # extensions = set([".bin", ".ckpt"])
+    # for filename in filenames:
+    #     prefix, ext = os.path.splitext(filename)
+    #     if ext in extensions:
+    #         pt_filename = hf_hub_download(
+    #             model_id, revision=revision, filename=filename, token=token, cache_dir=folder
+    #         )
+    #         dirname, raw_filename = os.path.split(filename)
+    #         if raw_filename == "pytorch_model.bin":
+    #             # XXX: This is a special case to handle `transformers` and the
+    #             # `transformers` part of the model which is actually loaded by `transformers`.
+    #             sf_in_repo = os.path.join(dirname, "model.safetensors")
+    #         else:
+    #             sf_in_repo = f"{prefix}.safetensors"
+    #         sf_filename = os.path.join(folder, sf_in_repo)
+    #         try:
+    #             convert_file(pt_filename, sf_filename, discard_names=[])
+    #             operations.append(CommitOperationAdd(path_in_repo=sf_in_repo, path_or_fileobj=sf_filename))
+    #         except Exception as e:
+    #             errors.append((pt_filename, e))
+    return operations, errors
+
+def quantize(
+    api: "HfApi", model_id: str, revision: Optional[str] = None, force: bool = False
+) -> Tuple["CommitInfo", List[Tuple[str, "Exception"]]]:
+    pr_title = "Adding `CoreML` variant of this model"
+    # info = api.model_info(model_id, revision=revision)
+    # filenames = set(s.rfilename for s in info.siblings)
+
+    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:
+            operations = None
+            pr = None
+
+            operations, errors = convert_generic(model_id, revision=revision, folder=folder, filenames={"pytorch_model.bin"}, token=api.token)
+
+            new_pr = api.create_commit(
+                repo_id=model_id,
+                revision=revision,
+                operations=operations,
+                commit_message=pr_title,
+                commit_description="Add CoreML variant of this model",
+                create_pr=True,
+            )
+            print(f"Pr created at {new_pr.pr_url}")
+        finally:
+            shutil.rmtree(folder)
+        return new_pr, errors
+
+
+if __name__ == "__main__":
+    DESCRIPTION = """
+    Simple utility tool to convert automatically some weights on the hub to `safetensors` format.
+    It is PyTorch exclusive for now.
+    It works by downloading the weights (PT), converting them locally, and uploading them back
+    as a PR on the hub.
+    """
+    parser = argparse.ArgumentParser(description=DESCRIPTION)
+    parser.add_argument(
+        "model_id",
+        type=str,
+        help="The name of the model on the hub to convert. E.g. `gpt2` or `facebook/wav2vec2-base-960h`",
+    )
+    parser.add_argument(
+        "--revision",
+        type=str,
+        help="The revision to convert",
+    )
+    parser.add_argument(
+        "--force",
+        action="store_true",
+        help="Create the PR even if it already exists of if the model was already converted.",
+    )
+    parser.add_argument(
+        "-y",
+        action="store_true",
+        help="Ignore safety prompt",
+    )
+    args = parser.parse_args()
+    model_id = args.model_id
+    api = HfApi()
+    if args.y:
+        txt = "y"
+    else:
+        txt = input(
+            "This conversion script will unpickle a pickled file, which is inherently unsafe. If you do not trust this file, we invite you to use"
+            " https://huggingface.co/spaces/safetensors/convert or google colab or other hosted solution to avoid potential issues with this file."
+            " Continue [Y/n] ?"
+        )
+    if txt.lower() in {"", "y"}:
+        commit_info, errors = convert(api, model_id, revision=args.revision, force=args.force)
+        string = f"""
+### Success 🔥
+Yay! This model was successfully converted and a PR was open using your token, here:
+[{commit_info.pr_url}]({commit_info.pr_url})
+        """
+        if errors:
+            string += "\nErrors during conversion:\n"
+            string += "\n".join(
+                f"Error while converting {filename}: {e}, skipped conversion" for filename, e in errors
+            )
+        print(string)
+    else:
+        print(f"Answer was `{txt}` aborting.")

requirements.txt

@@ -0,0 +1,105 @@
+accelerate==0.24.1
+aiofiles==23.2.1
+altair==5.3.0
+annotated-types==0.6.0
+antlr4-python3-runtime==4.9.3
+anyio==4.3.0
+attrs==23.1.0
+cattrs==23.1.2
+certifi==2023.11.17
+charset-normalizer==3.3.2
+click==8.1.7
+contourpy==1.1.1
+coremltools==7.1
+cycler==0.12.1
+diffusers==0.22.3
+exceptiongroup==1.1.3
+fastapi==0.110.2
+ffmpy==0.3.2
+filelock==3.13.1
+fonttools==4.44.0
+fsspec==2023.12.2
+gradio==4.27.0
+gradio_client==0.15.1
+h11==0.14.0
+httpcore==1.0.5
+httpx==0.27.0
+huggingface-hub==0.22.2
+idna==3.6
+importlib-metadata==6.8.0
+importlib-resources==6.1.1
+iniconfig==2.0.0
+invisible-watermark==0.2.0
+Jinja2==3.1.2
+joblib==1.3.2
+jsonschema==4.21.1
+jsonschema-specifications==2023.12.1
+kiwisolver==1.4.5
+lightning-utilities==0.11.2
+markdown-it-py==3.0.0
+MarkupSafe==2.1.3
+matplotlib==3.7.3
+mdurl==0.1.2
+mpmath==1.3.0
+networkx==3.1
+ninja==1.11.1.1
+numpy==1.23.5
+omegaconf==2.3.0
+opencv-python==4.8.1.78
+orjson==3.10.1
+packaging==23.2
+pandas==2.0.3
+Pillow==10.1.0
+pkgutil_resolve_name==1.3.10
+pluggy==1.3.0
+protobuf==3.20.3
+psutil==5.9.6
+pyaml==23.9.7
+pydantic==2.7.0
+pydantic_core==2.18.1
+pydub==0.25.1
+Pygments==2.17.2
+pyparsing==3.1.1
+pytest==7.4.3
+-e git+https://github.com/apple/ml-stable-diffusion.git@f7d2f5e9fb0681b15770943e492bf2f6dc3414f3#egg=python_coreml_stable_diffusion
+python-dateutil==2.8.2
+python-multipart==0.0.9
+pytz==2024.1
+PyWavelets==1.4.1
+PyYAML==6.0.1
+quanto==0.1.0
+referencing==0.34.0
+regex==2023.12.25
+requests==2.31.0
+rich==13.7.1
+rpds-py==0.18.0
+ruff==0.4.1
+safetensors==0.4.1
+scikit-learn==1.1.2
+scipy==1.10.1
+semantic-version==2.10.0
+sentencepiece==0.2.0
+shellingham==1.5.4
+six==1.16.0
+sniffio==1.3.1
+starlette==0.37.2
+sympy==1.12
+threadpoolctl==3.2.0
+tokenizers==0.14.1
+tomli==2.0.1
+tomlkit==0.12.0
+toolz==0.12.1
+torch==2.1.0
+torchao==0.1
+torchmetrics==1.3.2
+tqdm==4.66.1
+transformers==4.34.1
+typer==0.12.3
+typing_extensions==4.9.0
+tzdata==2024.1
+urllib3==2.1.0
+uvicorn==0.29.0
+websockets==11.0.3
+zipp==3.17.0
+setuptools_rust
+pytorch_lightning