Additional args

app.py

@@ -14,13 +14,7 @@ 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, model_version: str, is_private: bool, token: Optional[str] = None) -> str:
+def run(model_id: str, model_version: str, additional_args: str, is_private: bool, token: Optional[str] = None) -> str:
     if model_id == "":
         return """
         ### Invalid input 🐞
@@ -28,38 +22,19 @@ def run(model_id: str, model_version: str, is_private: bool, token: Optional[str
         Please fill a token and model_id.
         """
     login(token=token)
-    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, model_version=model_version)
-        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}
-        """
+    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)
+
+    quantize.quantize(api=api, model_id=model_id, model_version=model_version, additional_args=additional_args)
 
 
 DESCRIPTION = """
@@ -85,6 +60,7 @@ with gr.Blocks(title=title) as demo:
         with gr.Column() as c:
             model_id = gr.Text(max_lines=1, label="model_id", value="jblalock30/coreml")
             model_version = gr.Text(max_lines=1, label="model_version", value="stabilityai/sd-turbo")
+            additional_args = gr.Text(max_lines=1, label="additional_args", value="--quantize-nbits 2 --convert-unet --convert-text-encoder --convert-vae-decoder --chunk-unet --attention-implementation ORIGINAL")
             is_private = gr.Checkbox(label="Private model")
             token = token_text()
             with gr.Row() as c:
@@ -92,9 +68,9 @@ with gr.Blocks(title=title) as demo:
                 submit = gr.Button("Submit", variant="primary")
 
         with gr.Column() as d:
-            output = gr.Markdown(value="hi")
+            output = gr.Markdown()
 
     is_private.change(lambda s: token_text(s), inputs=is_private, outputs=token)
-    submit.click(run, inputs=[model_id, model_version, is_private, token], outputs=output, concurrency_limit=1)
+    submit.click(run, inputs=[model_id, model_version, additional_args, is_private, token], outputs=output, concurrency_limit=1)
 
 demo.queue(max_size=10).launch(show_api=True)

quantize.py

@@ -16,90 +16,37 @@ from safetensors.torch import _find_shared_tensors, _is_complete, load_file, sav
 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], model_version: str
+    model_id: str, folder: str, token: Optional[str], model_version: str, additional_args: str
 ) -> ConversionResult:
-    operations = []
-    errors = []
+    
+    command = ["python3", "-m" , "python_coreml_stable_diffusion.torch2coreml", "--model-version", model_version, "-o", folder]
+    command.extend(additional_args.split(" "))
 
-    # 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")
 
-    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"])
-    subprocess.run(["python3", "-m" , "python_coreml_stable_diffusion.torch2coreml", "--model-version", model_version, "-o", folder, "--convert-unet", "--convert-text-encoder", "--convert-vae-decoder", "--chunk-unet", "--attention-implementation", "ORIGINAL", "--quantize-nbits", "2"])
-    # with open(f'{folder}/newfile.txt', 'w') as f:
-    #     f.write('Hello, World!')
+    subprocess.run(command)
 
     print("Done")
 
     api = HfApi(token=token)
     api.upload_folder(
         folder_path=folder,
-        repo_id="jblalock30/coreml",
+        repo_id=model_id,
         path_in_repo="models",
         repo_type="model",
     )
-    # for filename in os.listdir(folder):
-    #     print(filename)
-
-    #     operations.append(CommitOperationAdd(path_in_repo=f'models/{filename}', path_or_fileobj=f'{folder}/{filename}'))
-
-    # 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, model_version: 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)
-
+    api: "HfApi", model_id: str, model_version: str, additional_args: str
+) -> None:
+    
     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, model_version=model_version, 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}")
+            convert_generic(model_id, folder, token=api.token, model_version=model_version, additional_args=additional_args)
         finally:
             shutil.rmtree(folder)
-        return new_pr, errors
 
 
 if __name__ == "__main__":