Move to utils

app.py

@@ -475,7 +475,9 @@ if __name__ == "__main__":
     args = parser.parse_args()
 
     # download to hub cache
-    download_repo_files()
+    for repo_id, allow_patterns in Config.DOWNLOAD_FILES.items():
+        print(f"Downloading {repo_id}...")
+        download_repo_files(repo_id, allow_patterns, token=Config.HF_TOKEN)
 
     # https://www.gradio.app/docs/gradio/interface#interface-queue
     demo.queue().launch(

lib/__init__.py

@@ -1,9 +1,8 @@
 from .config import Config
-from .download import download_repo_files
 from .inference import generate
 from .loader import Loader
 from .upscaler import RealESRGAN
-from .utils import async_call, load_json, read_file
+from .utils import async_call, download_repo_files, load_json, read_file
 
 __all__ = [
     "Config",

lib/config.py

@@ -1,3 +1,4 @@
+import os
 from types import SimpleNamespace
 
 from diffusers import (
@@ -12,6 +13,25 @@ from diffusers import (
 )
 
 Config = SimpleNamespace(
+    HF_TOKEN=os.environ.get("HF_TOKEN", None),
+    DOWNLOAD_FILES={
+        "Lykon/dreamshaper-8": [
+            "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",
+        ],
+    },
     MONO_FONTS=["monospace"],
     SANS_FONTS=[
         "sans-serif",

lib/download.py

@@ -1,38 +0,0 @@
-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,
-    )

lib/utils.py

@@ -5,6 +5,7 @@ from typing import Callable, TypeVar
 
 import anyio
 from anyio import Semaphore
+from huggingface_hub._snapshot_download import snapshot_download
 from typing_extensions import ParamSpec
 
 T = TypeVar("T")
@@ -26,6 +27,17 @@ def read_file(path: str) -> str:
         return file.read()
 
 
+def download_repo_files(repo_id, allow_patterns, token=None):
+    return snapshot_download(
+        repo_id=repo_id,
+        repo_type="model",
+        revision="main",
+        token=token,
+        allow_patterns=allow_patterns,
+        ignore_patterns=None,
+    )
+
+
 # like the original but supports args and kwargs instead of a dict
 # https://github.com/huggingface/huggingface-inference-toolkit/blob/0.2.0/src/huggingface_inference_toolkit/async_utils.py
 async def async_call(fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T: