Add split_einsum example to download_model.py

scripts/README.md

@@ -2,17 +2,21 @@
 
 ## Download models
 
-Modify `scripts/download_model.py` and run it: `python scripts/download_model.py`
+Modify `scripts/download_model.py` to choose the model you'd like to download.
+Then run it from the command line: `python scripts/download_model.py`
+
+Check the `models/<output>` folder to see the model. If it is a https://huggingface.co/coreml model (models converted to Core ML), then unzip it first.
 
 ## 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
+# MODEL=coreml-stable-diffusion-v1-5_original_compiled
+MODEL=coreml-stable-diffusion-2-1-base_original
+# MODEL=coreml-stable-diffusion-2-1-base_split_einsum
 # COMPUTE_UNITS=all # "split_einsum" models
 COMPUTE_UNITS=cpuAndGPU # on "original" models
+OUTPUT_PATH=output_images/$MODEL
 mkdir -p $OUTPUT_PATH
 
 PROMPT="a photograph of an astronaut riding on a horse"

scripts/download_model.py

@@ -4,8 +4,8 @@ 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-v1-5"
 # repo_id = "apple/coreml-stable-diffusion-2-base"
 
 # For Swift
@@ -16,20 +16,27 @@ import shutil
 # From coreml: https://huggingface.co/coreml
 repo_id = "coreml/coreml-stable-diffusion-2-1-base"
 variant = "original"
+# variant = "split_einsum"
 
 
 def download_model(repo_id, variant, output_dir):
-    destination = Path(output_dir) / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
+    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 = 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")
+    cache_folder = Path(output_dir) / repo_folder_name(
+        repo_id=repo_id, repo_type="model"
+    )
     shutil.rmtree(cache_folder)
     return destination