handle loading invalid model or pipeline

modules/processing.py

@@ -279,7 +279,7 @@ class Processed:
         self.batch_size = p.batch_size
         self.restore_faces = p.restore_faces
         self.face_restoration_model = shared.opts.face_restoration_model if p.restore_faces else None
-        self.sd_model_hash = shared.sd_model.sd_model_hash
+        self.sd_model_hash = getattr(shared.sd_model, 'sd_model_hash', '')
         self.seed_resize_from_w = p.seed_resize_from_w
         self.seed_resize_from_h = p.seed_resize_from_h
         self.denoising_strength = p.denoising_strength
@@ -446,6 +446,8 @@ def fix_seed(p):
 
 
 def create_infotext(p: StableDiffusionProcessing, all_prompts, all_seeds, all_subseeds, comments=None, iteration=0, position_in_batch=0, index=None, all_negative_prompts=None):
+    if not hasattr(shared.sd_model, 'sd_checkpoint_info'):
+        return ''
     if index is None:
         index = position_in_batch + iteration * p.batch_size
     if all_negative_prompts is None:
@@ -533,6 +535,8 @@ def print_profile(profile, msg: str):
 
 
 def process_images(p: StableDiffusionProcessing) -> Processed:
+    if not hasattr(p.sd_model, 'sd_checkpoint_info'):
+        return None
     stored_opts = {}
     for k in p.override_settings.keys():
         stored_opts[k] = shared.opts.data.get(k, None)
@@ -584,7 +588,7 @@ def validate_sample(sample):
     try:
         sample = sample.astype(np.uint8)
         return sample
-    except (Warning, Exception) as e:
+    except (Exception, Warning, RuntimeWarning) as e:
         shared.log.error(f'Failed to validate sample values: {e}')
         ok = False
     if not ok:
@@ -592,7 +596,7 @@ def validate_sample(sample):
             sample = np.nan_to_num(sample, nan=0, posinf=255, neginf=0)
             sample = sample.astype(np.uint8)
             shared.log.debug('Corrected sample values')
-        except (Warning, Exception) as e:
+        except (Exception, Warning, RuntimeWarning) as e:
             shared.log.error(f'Failed to correct sample values: {e}')
             sample = np.zeros_like(sample)
             sample = sample.astype(np.uint8)

modules/sd_models.py

@@ -378,6 +378,7 @@ def get_checkpoint_state_dict(checkpoint_info: CheckpointInfo, timer):
 
 
 def load_model_weights(model: torch.nn.Module, checkpoint_info: CheckpointInfo, state_dict, timer):
+    _pipeline, _model_type = detect_pipeline(checkpoint_info.path, 'model')
     shared.log.debug(f'Model weights loading: {memory_stats()}')
     timer.record("hash")
     if model_data.sd_dict == 'None':
@@ -388,6 +389,7 @@ def load_model_weights(model: torch.nn.Module, checkpoint_info: CheckpointInfo,
         model.load_state_dict(state_dict, strict=False)
     except Exception as e:
         shared.log.error(f'Error loading model weights: {checkpoint_info.filename} {e}')
+        return False
     del state_dict
     timer.record("apply")
     if shared.opts.sd_checkpoint_cache > 0:
@@ -427,6 +429,7 @@ def load_model_weights(model: torch.nn.Module, checkpoint_info: CheckpointInfo,
     vae_file, vae_source = sd_vae.resolve_vae(checkpoint_info.filename)
     sd_vae.load_vae(model, vae_file, vae_source)
     timer.record("vae")
+    return True
 
 
 def enable_midas_autodownload():
@@ -554,14 +557,18 @@ def detect_pipeline(f: str, op: str = 'model'):
             size = round(os.path.getsize(f) / 1024 / 1024 / 1024, 2)
             if size < 1:
                 shared.log.warning(f'Model size smaller than expected: {f} size={size} GB')
-            elif size < 5:
+            elif size < 5.5: # maximum size of sd1.5 fp32 unpruned is 5.3GB
                 guess = 'Stable Diffusion'
-            elif size < 6:
+            elif size < 6: # sdxl refiner is 5.7gb
+                if shared.backend == shared.Backend.ORIGINAL:
+                    shared.log.warning(f'Model detected as SD-XL refiner model, but attempting to load using backend=original: {f} size={size} GB')
                 if op == 'model':
                     shared.log.warning(f'Model detected as SD-XL refiner model, but attempting to load a base model: {f} size={size} GB')
                 else:
                     guess = 'Stable Diffusion XL'
             elif size < 7:
+                if shared.backend == shared.Backend.ORIGINAL:
+                    shared.log.warning(f'Model detected as SD-XL base model, but attempting to load using backend=original: {f} size={size} GB')
                 if op == 'refiner':
                     shared.log.warning(f'Model size matches SD-XL base model, but attempting to load a refiner model: {f} size={size} GB')
                 else:
@@ -965,9 +972,16 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, timer=None,
     shared.log.debug(f"Model created from config: {checkpoint_config}")
     sd_model.used_config = checkpoint_config
     timer.record("create")
-    load_model_weights(sd_model, checkpoint_info, state_dict, timer)
+    ok = load_model_weights(sd_model, checkpoint_info, state_dict, timer)
+    if not ok:
+        model_data.sd_model = sd_model
+        current_checkpoint_info = None
+        unload_model_weights(op=op)
+        shared.log.debug(f'Model weights unloaded: {memory_stats()}')
+        return
+    else:
+        shared.log.debug(f'Model weights loaded: {memory_stats()}')
     timer.record("load")
-    shared.log.debug(f'Model weights loaded: {memory_stats()}')
     if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
         lowvram.setup_for_low_vram(sd_model, shared.cmd_opts.medvram)
     else:
@@ -991,6 +1005,7 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, timer=None,
     devices.torch_gc(force=True)
     shared.log.info(f'Model load finished: {memory_stats()} cached={len(checkpoints_loaded.keys())}')
 
+
 def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model'):
     load_dict = shared.opts.sd_model_dict != model_data.sd_dict
     from modules import lowvram, sd_hijack

modules/shared.py

@@ -814,7 +814,6 @@ opts.data['uni_pc_lower_order_final'] = opts.schedulers_use_loworder
 opts.data['uni_pc_order'] = opts.schedulers_solver_order
 log.info(f'Engine: backend={backend}')
 
-
 prompt_styles = modules.styles.StyleDatabase(opts.styles_dir)
 cmd_opts.disable_extension_access = (cmd_opts.share or cmd_opts.listen or (cmd_opts.server_name or False)) and not cmd_opts.insecure
 devices.device, devices.device_interrogate, devices.device_gfpgan, devices.device_esrgan, devices.device_codeformer = (devices.cpu if any(y in cmd_opts.use_cpu for y in [x, 'all']) else devices.get_optimal_device() for x in ['sd', 'interrogate', 'gfpgan', 'esrgan', 'codeformer'])

modules/txt2img.py

@@ -64,6 +64,8 @@ def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, step
     if processed is None:
         processed = processing.process_images(p)
     p.close()
+    if processed is None:
+        return [], '', '', 'Error: processing failed'
     generation_info_js = processed.js()
     shared.log.debug(f'Processed: {len(processed.images)} Memory: {memory_stats()} txt')
     return processed.images, generation_info_js, processed.info, plaintext_to_html(processed.comments)

modules/ui_extra_networks.py

@@ -15,13 +15,15 @@ from modules.ui_components import ToolButton
 
 extra_pages = []
 allowed_dirs = set()
-dir_cache = {}
+dir_cache = {} # key=path, value=(mtime, listdir(path))
 
 refresh_symbol = '\U0001f504'  # 🔄
 close_symbol = '\U0000274C'  # ❌
 
 
 def listdir(path):
+    if not os.path.exists(path):
+        return []
     if path in dir_cache and os.path.getmtime(path) == dir_cache[path][0]:
         return dir_cache[path][1]
     else:

webui.py

@@ -15,12 +15,12 @@ from modules import timer, errors, paths # pylint: disable=unused-import
 startup_timer = timer.Timer()
 local_url = None
 
-errors.log.debug('Loading Torch')
 import torch # pylint: disable=C0411
 try:
     import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import
 except Exception:
     pass
+errors.log.debug(f'Loaded Torch=={torch.__version__}')
 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 import torchvision # pylint: disable=W0611,C0411
 import pytorch_lightning # pytorch_lightning should be imported after torch, but it re-enables warnings on import so import once to disable them # pylint: disable=W0611,C0411
@@ -34,12 +34,16 @@ warnings.filterwarnings(action="ignore", category=FutureWarning)
 warnings.filterwarnings(action="ignore", category=UserWarning, module="torchvision")
 startup_timer.record("torch")
 
-errors.log.debug('Loading Gradio')
 from fastapi import FastAPI # pylint: disable=W0611,C0411
 import gradio # pylint: disable=W0611,C0411
+errors.log.debug(f'Loaded Gradio=={gradio.__version__}')
 startup_timer.record("gradio")
 errors.install([gradio])
 
+import diffusers # pylint: disable=W0611,C0411
+errors.log.debug(f'Loaded Diffusers=={diffusers.__version__}')
+startup_timer.record("diffusers")
+
 errors.log.debug('Loading Modules')
 from installer import log, setup_logging, git_commit
 import ldm.modules.encoders.modules # pylint: disable=W0611,C0411,E0401
@@ -279,7 +283,7 @@ def start_ui():
         auth=[tuple(cred.split(':')) for cred in gradio_auth_creds] if gradio_auth_creds else None,
         prevent_thread_lock=True,
         max_threads=64,
-        show_api=True,
+        show_api=False,
         quiet=True,
         favicon_path='html/logo.ico',
         allowed_paths=[os.path.dirname(__file__), cmd_opts.data_dir],
@@ -304,6 +308,8 @@ def start_ui():
 
     modules.progress.setup_progress_api(app)
     create_api(app)
+    startup_timer.record("api")
+
     ui_extra_networks.add_pages_to_demo(app)
 
     modules.script_callbacks.app_started_callback(shared.demo, app)