Terminal progress bar improvements

app.py

@@ -1,23 +1,10 @@
 import argparse
 import json
-import os
 import random
-from warnings import filterwarnings
 
 import gradio as gr
-from diffusers.utils import logging as diffusers_logging
-from transformers import logging as transformers_logging
 
-from lib import Config, async_call, download_repo_files, generate, read_file
-
-filterwarnings("ignore", category=FutureWarning, module="diffusers")
-filterwarnings("ignore", category=FutureWarning, module="transformers")
-
-diffusers_logging.set_verbosity_error()
-transformers_logging.set_verbosity_error()
-
-diffusers_logging.disable_progress_bar()
-transformers_logging.disable_progress_bar()
+from lib import Config, async_call, disable_progress_bars, download_repo_files, generate, read_file
 
 # the CSS `content` attribute expects a string so we need to wrap the number in quotes
 refresh_seed_js = """
@@ -337,7 +324,7 @@ if __name__ == "__main__":
     parser.add_argument("-p", "--port", type=int, metavar="INT", default=7860)
     args = parser.parse_args()
 
-    # download to hub cache
+    disable_progress_bars()
     for repo_id, allow_patterns in Config.HF_MODELS.items():
         download_repo_files(repo_id, allow_patterns, token=Config.HF_TOKEN)
 

lib/__init__.py

@@ -2,15 +2,25 @@ from .config import Config
 from .inference import generate
 from .loader import Loader
 from .upscaler import RealESRGAN
-from .utils import async_call, download_civit_file, download_repo_files, load_json, read_file
+from .utils import (
+    async_call,
+    disable_progress_bars,
+    download_civit_file,
+    download_repo_files,
+    enable_progress_bars,
+    load_json,
+    read_file,
+)
 
 __all__ = [
     "Config",
     "Loader",
     "RealESRGAN",
     "async_call",
+    "disable_progress_bars",
     "download_civit_file",
     "download_repo_files",
+    "enable_progress_bars",
     "generate",
     "load_json",
     "read_file",

lib/config.py

@@ -1,6 +1,8 @@
 import os
 from importlib import import_module
+from importlib.util import find_spec
 from types import SimpleNamespace
+from warnings import filterwarnings
 
 from diffusers import (
     DDIMScheduler,
@@ -11,9 +13,21 @@ from diffusers import (
     StableDiffusionXLImg2ImgPipeline,
     StableDiffusionXLPipeline,
 )
+from diffusers.utils import logging as diffusers_logging
+from transformers import logging as transformers_logging
 
 # improved GPU handling and progress bars; set before importing spaces
-os.environ["ZEROGPU_V2"] = "true"
+os.environ["ZEROGPU_V2"] = "1"
+
+# use Rust downloader
+if find_spec("hf_transfer"):
+    os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+
+filterwarnings("ignore", category=FutureWarning, module="diffusers")
+filterwarnings("ignore", category=FutureWarning, module="transformers")
+
+diffusers_logging.set_verbosity_error()
+transformers_logging.set_verbosity_error()
 
 _sdxl_refiner_files = [
     "scheduler/scheduler_config.json",

lib/utils.py

@@ -7,7 +7,10 @@ from typing import Callable, TypeVar
 import anyio
 import httpx
 from anyio import Semaphore
+from diffusers.utils import logging as diffusers_logging
 from huggingface_hub._snapshot_download import snapshot_download
+from huggingface_hub.utils import are_progress_bars_disabled
+from transformers import logging as transformers_logging
 from typing_extensions import ParamSpec
 
 T = TypeVar("T")
@@ -29,8 +32,21 @@ def read_file(path: str) -> str:
         return file.read()
 
 
+def disable_progress_bars():
+    transformers_logging.disable_progress_bar()
+    diffusers_logging.disable_progress_bar()
+
+
+def enable_progress_bars():
+    # warns if `HF_HUB_DISABLE_PROGRESS_BARS` env var is not None
+    transformers_logging.enable_progress_bar()
+    diffusers_logging.enable_progress_bar()
+
+
 def download_repo_files(repo_id, allow_patterns, token=None):
-    return snapshot_download(
+    was_disabled = are_progress_bars_disabled()
+    enable_progress_bars()
+    snapshot_path = snapshot_download(
         repo_id=repo_id,
         repo_type="model",
         revision="main",
@@ -38,6 +54,9 @@ def download_repo_files(repo_id, allow_patterns, token=None):
         allow_patterns=allow_patterns,
         ignore_patterns=None,
     )
+    if was_disabled:
+        disable_progress_bars()
+    return snapshot_path
 
 
 def download_civit_file(lora_id, version_id, file_path=".", token=None):