Ensure files are downloaded

README.md

@@ -21,6 +21,7 @@ models:
 - h94/IP-Adapter
 - Linaqruf/anything-v3-1
 - Lykon/dreamshaper-8
+- madebyollin/taesd
 - prompthero/openjourney-v4
 - SG161222/Realistic_Vision_V5.1_noVAE
 - XpucT/Deliberate
@@ -45,7 +46,10 @@ preload_from_hub:
   anything-v3-2.safetensors
 - >-
   Lykon/dreamshaper-8
-  text_encoder/model.fp16.safetensors,unet/diffusion_pytorch_model.fp16.safetensors,vae/diffusion_pytorch_model.fp16.safetensors,model_index.json
+  feature_extractor/preprocessor_config.json,safety_checker/config.json,scheduler/scheduler_config.json,text_encoder/config.json,text_encoder/model.fp16.safetensors,tokenizer/merges.txt,tokenizer/special_tokens_map.json,tokenizer/tokenizer_config.json,tokenizer/vocab.json,unet/config.json,unet/diffusion_pytorch_model.fp16.safetensors,vae/config.json,vae/diffusion_pytorch_model.fp16.safetensors,model_index.json
+- >-
+  madebyollin/taesd
+  config.json,diffusion_pytorch_model.safetensors
 - >-
   prompthero/openjourney-v4
   openjourney-v4.ckpt

app.py

@@ -4,7 +4,7 @@ import random
 
 import gradio as gr
 
-from lib import Config, async_call, generate
+from lib import Config, async_call, download_repo_files, generate
 
 # the CSS `content` attribute expects a string so we need to wrap the number in quotes
 refresh_seed_js = """
@@ -476,6 +476,9 @@ if __name__ == "__main__":
     parser.add_argument("-p", "--port", type=int, metavar="INT", default=7860)
     args = parser.parse_args()
 
+    # download to hub cache
+    download_repo_files()
+
     # https://www.gradio.app/docs/gradio/interface#interface-queue
     demo.queue().launch(
         server_name=args.server,

lib/__init__.py

@@ -1,6 +1,7 @@
 from .config import Config
+from .download import download_repo_files
 from .inference import async_call, generate
 from .loader import Loader
 from .upscaler import RealESRGAN
 
-__all__ = ["Config", "Loader", "RealESRGAN", "async_call", "generate"]
+__all__ = ["Config", "Loader", "RealESRGAN", "async_call", "download_repo_files", "generate"]

lib/download.py

@@ -0,0 +1,38 @@
+import os
+
+from huggingface_hub._snapshot_download import snapshot_download
+
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
+
+SPACES_ZERO_GPU = os.environ.get("SPACES_ZERO_GPU", "").lower() == "true"
+
+REPO = "Lykon/dreamshaper-8"
+
+FILES = [
+    "feature_extractor/preprocessor_config.json",
+    "safety_checker/config.json",
+    "scheduler/scheduler_config.json",
+    "text_encoder/config.json",
+    "text_encoder/model.fp16.safetensors",
+    "tokenizer/merges.txt",
+    "tokenizer/special_tokens_map.json",
+    "tokenizer/tokenizer_config.json",
+    "tokenizer/vocab.json",
+    "unet/config.json",
+    "unet/diffusion_pytorch_model.fp16.safetensors",
+    "vae/config.json",
+    "vae/diffusion_pytorch_model.fp16.safetensors",
+    "model_index.json",
+]
+
+
+def download_repo_files():
+    global REPO, FILES
+    return snapshot_download(
+        repo_id=REPO,
+        repo_type="model",
+        revision="main",
+        token=HF_TOKEN,
+        allow_patterns=FILES,
+        ignore_patterns=None,
+    )