Terminal progress bar improvements

app.py

@@ -2,22 +2,18 @@ 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_civit_file, 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_civit_file,
+    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 = """
@@ -567,7 +563,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

@@ -3,7 +3,15 @@ from .inference import generate
 from .loader import Loader
 from .logger import Logger, log_fn
 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",
@@ -11,8 +19,10 @@ __all__ = [
     "Logger",
     "RealESRGAN",
     "async_call",
+    "disable_progress_bars",
     "download_civit_file",
     "download_repo_files",
+    "enable_progress_bars",
     "generate",
     "load_json",
     "log_fn",

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,33 +13,54 @@ from diffusers import (
     PNDMScheduler,
     UniPCMultistepScheduler,
 )
+from diffusers.utils import logging as diffusers_logging
+from transformers import logging as transformers_logging
 
 from .pipelines import CustomStableDiffusionImg2ImgPipeline, CustomStableDiffusionPipeline
 
 # improved GPU handling and progress bars; set before importing spaces
-os.environ["ZEROGPU_V2"] = "true"
+os.environ["ZEROGPU_V2"] = "1"
+
+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()
+
+_sd_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",
+]
 
 Config = SimpleNamespace(
     HF_TOKEN=os.environ.get("HF_TOKEN", None),
     CIVIT_TOKEN=os.environ.get("CIVIT_TOKEN", None),
     ZERO_GPU=import_module("spaces").config.Config.zero_gpu,
     HF_MODELS={
-        "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",
-        ],
+        # downloaded on startup
+        "Lykon/dreamshaper-8": [*_sd_files],
+        "Comfy-Org/stable-diffusion-v1-5-archive": ["v1-5-pruned-emaonly-fp16.safetensors"],
+        "cyberdelia/CyberRealistic": ["CyberRealistic_V5_FP16.safetensors"],
+        "fluently/Fluently-v4": ["Fluently-v4.safetensors"],
+        "Linaqruf/anything-v3-1": ["anything-v3-2.safetensors"],
+        "prompthero/openjourney-v4": ["openjourney-v4.ckpt"],
+        "SG161222/Realistic_Vision_V5.1_noVAE": ["Realistic_Vision_V5.1_fp16-no-ema.safetensors"],
+        "XpucT/Deliberate": ["Deliberate_v6.safetensors"],
     },
     CIVIT_LORAS={
         # https://civitai.com/models/411088?modelVersionId=486099

lib/utils.py

@@ -7,15 +7,22 @@ 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
 
+from .logger import Logger
+
 T = TypeVar("T")
 P = ParamSpec("P")
 
 MAX_CONCURRENT_THREADS = 1
 MAX_THREADS_GUARD = Semaphore(MAX_CONCURRENT_THREADS)
 
+log = Logger("utils")
+
 
 @functools.lru_cache()
 def load_json(path: str) -> dict:
@@ -29,8 +36,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 +58,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):
@@ -62,9 +85,9 @@ def download_civit_file(lora_id, version_id, file_path=".", token=None):
         with open(file, "wb") as f:
             f.write(response.content)
     except httpx.HTTPStatusError as e:
-        print(f"HTTPError: {e.response.status_code} {e.response.text}")
+        log.error(f"{e.response.status_code} {e.response.text}")
     except httpx.RequestError as e:
-        print(f"RequestError: {e}")
+        log.error(f"RequestError: {e}")
 
 
 # like the original but supports args and kwargs instead of a dict