From ec3cbb829c172187dd238346f08487538435152b Mon Sep 17 00:00:00 2001
From: Brian Sigafoos <sigafoos@gmail.com>
Date: Tue, 24 Jan 2023 17:49:06 -0500
Subject: [PATCH] Add scripts/download_model.py and README (#1)

---
 .gitignore                |  4 +++
 scripts/README.md         | 68 +++++++++++++++++++++++++++++++++++++++
 scripts/download_model.py | 37 +++++++++++++++++++++
 3 files changed, 109 insertions(+)
 create mode 100644 scripts/README.md
 create mode 100755 scripts/download_model.py

diff --git a/.gitignore b/.gitignore
index 5c57bef..13a46f1 100644
--- a/.gitignore
+++ b/.gitignore
@@ -142,3 +142,7 @@ dmypy.json
 
 # macOS filesystem
 *.DS_Store
+
+# Added
+models/
+output_images/
diff --git a/scripts/README.md b/scripts/README.md
new file mode 100644
index 0000000..e3dc0dc
--- /dev/null
+++ b/scripts/README.md
@@ -0,0 +1,68 @@
+# Download models and run inference
+
+## Download models
+
+Modify `scripts/download_model.py` and run it: `python scripts/download_model.py`
+
+## Run inference
+
+```shell
+MODEL=coreml-stable-diffusion-2-1-base_original
+# MODEL=coreml-stable-diffusion-v1-5_original_compiled
+# MODEL=coreml-stable-diffusion-v1-4_original_compiled
+OUTPUT_PATH=output_images/$MODEL
+# COMPUTE_UNITS=all # "split_einsum" models
+COMPUTE_UNITS=cpuAndGPU # on "original" models
+mkdir -p $OUTPUT_PATH
+
+PROMPT="a photograph of an astronaut riding on a horse"
+SEED=42 # 93 is the default
+echo "Generating \"$PROMPT\" on $MODEL with seed $SEED"
+time swift run StableDiffusionSample $PROMPT --resource-path models/$MODEL --compute-units $COMPUTE_UNITS --output-path $OUTPUT_PATH --seed $SEED
+```
+
+Available commands:
+
+```shell
+swift run StableDiffusionSample --help
+# USAGE: stable-diffusion-sample [<options>] <prompt>
+
+# ARGUMENTS:
+#   <prompt>                Input string prompt
+
+# OPTIONS:
+#   --negative-prompt <negative-prompt>
+#                           Input string negative prompt
+#   --resource-path <directory-path>
+#                           Path to stable diffusion resources. (default: ./)
+#         The resource directory should contain
+#          - *compiled* models: {TextEncoder,Unet,VAEDecoder}.mlmodelc
+#          - tokenizer info: vocab.json, merges.txt
+#   --image-count <image-count>
+#                           Number of images to sample / generate (default: 1)
+#   --step-count <step-count>
+#                           Number of diffusion steps to perform (default: 50)
+#   --save-every <save-every>
+#                           How often to save samples at intermediate steps (default: 0)
+#         Set to 0 to only save the final sample
+#   --output-path <output-path>
+#                           Output path (default: ./)
+#   --seed <seed>           Random seed (default: 93)
+#   --guidance-scale <guidance-scale>
+#                           Controls the influence of the text prompt on sampling process (0=random
+#                           images) (default: 7.5)
+#   --compute-units <compute-units>
+#                           Compute units to load model with
+#                           {all,cpuOnly,cpuAndGPU,cpuAndNeuralEngine} (default: all)
+#   --scheduler <scheduler> Scheduler to use, one of {pndm, dpmpp} (default: pndm)
+#   --disable-safety        Disable safety checking
+#   --reduce-memory         Reduce memory usage
+#   --version               Show the version.
+#   -h, --help              Show help information.
+```
+
+## References
+
+- https://huggingface.co/blog/diffusers-coreml
+- https://huggingface.co/apple
+- https://huggingface.co/coreml
diff --git a/scripts/download_model.py b/scripts/download_model.py
new file mode 100755
index 0000000..a75a873
--- /dev/null
+++ b/scripts/download_model.py
@@ -0,0 +1,37 @@
+from huggingface_hub import snapshot_download
+from huggingface_hub.file_download import repo_folder_name
+from pathlib import Path
+import shutil
+
+# From apple: https://huggingface.co/apple
+# repo_id = "apple/coreml-stable-diffusion-v1-5"
+# repo_id = "apple/coreml-stable-diffusion-v1-4"
+# repo_id = "apple/coreml-stable-diffusion-2-base"
+
+# For Swift
+# variant = "original/compiled"
+# For Python
+# variant = "original/packages"
+
+# From coreml: https://huggingface.co/coreml
+repo_id = "coreml/coreml-stable-diffusion-2-1-base"
+variant = "original"
+
+
+def download_model(repo_id, variant, output_dir):
+    destination = Path(output_dir) / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
+    if destination.exists():
+        raise Exception(f"Model already exists at {destination}")
+
+    # Download and copy without symlinks
+    downloaded = snapshot_download(repo_id, allow_patterns=f"{variant}/*", cache_dir=output_dir)
+    downloaded_bundle = Path(downloaded) / variant
+    shutil.copytree(downloaded_bundle, destination)
+
+    # Remove all downloaded files
+    cache_folder = Path(output_dir) / repo_folder_name(repo_id=repo_id, repo_type="model")
+    shutil.rmtree(cache_folder)
+    return destination
+
+model_path = download_model(repo_id, variant, output_dir="./models")
+print(f"Model downloaded at {model_path}")