From 2f3557f6ea4e93a2ab5a8756390171cbfb48039d Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Fri, 10 May 2024 10:08:01 -0400
Subject: [PATCH] patch diffusers loader

---
 modules/loader.py        | 5 +++++
 modules/merging/merge.py | 1 -
 modules/sd_models.py     | 2 +-
 3 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/modules/loader.py b/modules/loader.py
index 34039f4f4..64e617890 100644
--- a/modules/loader.py
+++ b/modules/loader.py
@@ -4,6 +4,7 @@ import sys
 import logging
 import warnings
 import urllib3
+from functools import partial
 from modules import timer, errors
 
 initialized = False
@@ -53,6 +54,10 @@ import pydantic # pylint: disable=W0611,C0411
 timer.startup.record("pydantic")
 
 import diffusers # pylint: disable=W0611,C0411
+import diffusers.loaders.single_file # pylint: disable=W0611,C0411
+logging.getLogger("diffusers.loaders.single_file").setLevel(logging.ERROR)
+from tqdm.rich import tqdm
+diffusers.loaders.single_file.logging.tqdm = partial(tqdm, unit='C')
 timer.startup.record("diffusers")
 
 def get_packages():
diff --git a/modules/merging/merge.py b/modules/merging/merge.py
index d1e6129bc..d6430033e 100644
--- a/modules/merging/merge.py
+++ b/modules/merging/merge.py
@@ -189,7 +189,6 @@ def simple_merge(
     threads: int = 4,
 ) -> Dict:
     futures = []
-    # with tqdm(thetas["model_a"].keys(), desc="Merge") as progress:
     import rich.progress as p
     with p.Progress(p.TextColumn('[cyan]{task.description}'), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TextColumn('[cyan]keys={task.fields[keys]}'), console=console) as progress:
         task = progress.add_task(description="Merging", total=len(thetas["model_a"].keys()), keys=len(thetas["model_a"].keys()))
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 219b28b50..b1841e6c7 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -1063,7 +1063,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
                     shared.log.error(f'Failed loading {op}: {checkpoint_info.path} auto={err1} diffusion={err2}')
                     return
         elif os.path.isfile(checkpoint_info.path) and checkpoint_info.path.lower().endswith('.safetensors'):
-            # diffusers_load_config["local_files_only"] = True
+            diffusers_load_config["local_files_only"] = True
             diffusers_load_config["extract_ema"] = shared.opts.diffusers_extract_ema
             if pipeline is None:
                 shared.log.error(f'Diffusers {op} pipeline not initialized: {shared.opts.diffusers_pipeline}')