reference load cleanup

Signed-off-by: vladmandic <mandic00@live.com>

modules/loader.py

@@ -145,6 +145,7 @@ except Exception as e:
     sys.exit(1)
 
 import huggingface_hub # pylint: disable=W0611,C0411
+logging.getLogger("huggingface_hub.file_download").setLevel(logging.ERROR)
 timer.startup.record("hfhub")
 
 try:
@@ -188,7 +189,7 @@ def get_packages():
 try:
     import math
     cores = os.cpu_count()
-    affinity = len(os.sched_getaffinity(0))
+    affinity = len(os.sched_getaffinity(0)) # pylint: disable=no-member
     threads = torch.get_num_threads()
     if threads < (affinity / 2):
         torch.set_num_threads(math.floor(affinity / 2))

modules/modelloader.py

@@ -79,7 +79,7 @@ def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config
         download_config["mirror"] = mirror
     if custom_pipeline is not None and len(custom_pipeline) > 0:
         download_config["custom_pipeline"] = custom_pipeline
-    shared.log.debug(f'Diffusers downloading: id="{hub_id}" args={download_config}')
+    shared.log.debug(f'HF download: id="{hub_id}" args={download_config}')
     token = token or shared.opts.huggingface_token
     if token is not None and len(token) > 2:
         hf_login(token)
@@ -94,20 +94,20 @@ def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config
         except Exception as e:
             err = e
             ok = False
-            debug(f'Diffusers download error: id="{hub_id}" {e}')
+            debug(f'HF download error: id="{hub_id}" {e}')
     if not ok and 'Repository Not Found' not in str(err):
         try:
             download_config.pop('load_connected_pipeline', None)
             download_config.pop('variant', None)
             pipeline_dir = hf.snapshot_download(hub_id, **download_config)
         except Exception as e:
-            debug(f'Diffusers download error: id="{hub_id}" {e}')
+            debug(f'HF download error: id="{hub_id}" {e}')
             if 'gated' in str(e):
-                shared.log.error(f'Diffusers download error: id="{hub_id}" model access requires login')
+                shared.log.error(f'HF download error: id="{hub_id}" model access requires login')
                 shared.state.end(jobid)
                 return None
     if pipeline_dir is None:
-        shared.log.error(f'Diffusers download error: id="{hub_id}" {err}')
+        shared.log.error(f'HF download error: id="{hub_id}" {err}')
         shared.state.end(jobid)
         return None
     try:

modules/sd_detect.py

@@ -179,7 +179,7 @@ def detect_pipeline(f: str, op: str = 'model'):
     pipeline = None
     if guess == 'Autodetect':
         try:
-            guess = 'Stable Diffusion XL' if 'XL' in f.upper() else 'Stable Diffusion' # set default guess
+            guess = 'Stable Diffusion XL' if ('XL' in f.upper() or 'SDNQ' in f.upper()) else 'Stable Diffusion' # set default guess
             guess = guess_by_size(f, guess)
             guess = guess_by_name(f, guess)
             guess, pipeline = guess_by_diffusers(f, guess)

modules/sd_models.py

@@ -425,6 +425,10 @@ def load_diffuser_folder(model_type, pipeline, checkpoint_info, diffusers_load_c
 
     try: #0 - using detected model type and pipeline
         if (model_type is not None) and (pipeline is not None):
+            if ('sdnq' in model_type.lower()) or ('sdnq' in checkpoint_info.path.lower()):
+                from modules import sdnq # pylint: disable=unused-import # register to diffusers and transformers
+                global allow_post_quant # pylint: disable=global-statement
+                allow_post_quant = False
             sd_model = pipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
             sd_model.model_type = sd_model.__class__.__name__
     except Exception as e:
@@ -466,7 +470,7 @@ def load_diffuser_folder(model_type, pipeline, checkpoint_info, diffusers_load_c
 
     try: # 3 - try basic pipeline just in case
         if err2 is not None:
-            sd_model = diffusers.StableDiffusionPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
+            sd_model = diffusers.StableDiffusionXLPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
             sd_model.model_type = sd_model.__class__.__name__
     except Exception as e:
         err3 = e  # ignore last error
@@ -712,16 +716,16 @@ def load_diffuser(checkpoint_info=None, op='model', revision=None): # pylint: di
                 allow_post_quant = False
                 model_type = model_type.replace(' SDNQ', '')
 
-        # load from hf folder-style
-        if sd_model is None:
-            if os.path.isdir(checkpoint_info.path) or checkpoint_info.type == 'huggingface' or checkpoint_info.type == 'transformer':
-                sd_model = load_diffuser_folder(model_type, pipeline, checkpoint_info, diffusers_load_config, op)
-
         # load from single-file
         if sd_model is None:
             if os.path.isfile(checkpoint_info.path) and checkpoint_info.path.lower().endswith('.safetensors'):
                 sd_model = load_diffuser_file(model_type, pipeline, checkpoint_info, diffusers_load_config, op)
 
+        # load from hf folder-style
+        if sd_model is None:
+            if os.path.isdir(checkpoint_info.path) or (checkpoint_info.type == 'huggingface') or (checkpoint_info.type == 'transformer') or (checkpoint_info.type == 'reference'):
+                sd_model = load_diffuser_folder(model_type, pipeline, checkpoint_info, diffusers_load_config, op)
+
         if sd_model is None:
             shared.log.error(f'Load {op}: name="{checkpoint_info.name if checkpoint_info is not None else None}" not loaded')
             return

modules/ui_settings.py

@@ -385,11 +385,19 @@ def create_quicksettings(interfaces):
         def reference_submit(model):
             if '@' not in model: # diffusers
                 loaded = modelloader.load_reference(model)
-                return model if loaded else shared.opts.sd_model_checkpoint
+                if loaded:
+                    shared.opts.sd_model_checkpoint = model
+                    sd_models.reload_model_weights(force=True)
+                    return model
+                return shared.opts.sd_model_checkpoint
             else: # civitai
                 model, url = model.split('@')
                 loaded = modelloader.load_civitai(model, url)
-                return loaded if loaded is not None else shared.opts.sd_model_checkpoint
+                if loaded is not None:
+                    shared.opts.sd_model_checkpoint = loaded.title
+                    sd_models.reload_model_weights(force=True)
+                    return loaded
+                return shared.opts.sd_model_checkpoint
 
         button_set_reference = gr.Button('Change reference', elem_id='change_reference', visible=False)
         button_set_reference.click(