Merge branch 'master' of github.com:AUTOMATIC1111/stable-diffusion-webui into gamepad

1 files changed, 238 insertions, 119 deletions

diff --git a/modules/sd_models.py b/modules/sd_models.py
index ecdd91c5..b2d48a51 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -2,7 +2,6 @@ import collections
 import os.path

 import sys

 import gc

-from collections import namedtuple

 import torch

 import re

 import safetensors.torch

@@ -13,17 +12,63 @@ import ldm.modules.midas as midas
 

 from ldm.util import instantiate_from_config

 

-from modules import shared, modelloader, devices, script_callbacks, sd_vae

+from modules import paths, shared, modelloader, devices, script_callbacks, sd_vae, sd_disable_initialization, errors, hashes, sd_models_config

 from modules.paths import models_path

-from modules.sd_hijack_inpainting import do_inpainting_hijack, should_hijack_inpainting

+from modules.sd_hijack_inpainting import do_inpainting_hijack

+from modules.timer import Timer

 

 model_dir = "Stable-diffusion"

-model_path = os.path.abspath(os.path.join(models_path, model_dir))

+model_path = os.path.abspath(os.path.join(paths.models_path, model_dir))

 

-CheckpointInfo = namedtuple("CheckpointInfo", ['filename', 'title', 'hash', 'model_name', 'config'])

 checkpoints_list = {}

+checkpoint_alisases = {}

 checkpoints_loaded = collections.OrderedDict()

 

+

+class CheckpointInfo:

+    def __init__(self, filename):

+        self.filename = filename

+        abspath = os.path.abspath(filename)

+

+        if shared.cmd_opts.ckpt_dir is not None and abspath.startswith(shared.cmd_opts.ckpt_dir):

+            name = abspath.replace(shared.cmd_opts.ckpt_dir, '')

+        elif abspath.startswith(model_path):

+            name = abspath.replace(model_path, '')

+        else:

+            name = os.path.basename(filename)

+

+        if name.startswith("\\") or name.startswith("/"):

+            name = name[1:]

+

+        self.name = name

+        self.model_name = os.path.splitext(name.replace("/", "_").replace("\\", "_"))[0]

+        self.hash = model_hash(filename)

+

+        self.sha256 = hashes.sha256_from_cache(self.filename, "checkpoint/" + name)

+        self.shorthash = self.sha256[0:10] if self.sha256 else None

+

+        self.title = name if self.shorthash is None else f'{name} [{self.shorthash}]'

+

+        self.ids = [self.hash, self.model_name, self.title, name, f'{name} [{self.hash}]'] + ([self.shorthash, self.sha256, f'{self.name} [{self.shorthash}]'] if self.shorthash else [])

+

+    def register(self):

+        checkpoints_list[self.title] = self

+        for id in self.ids:

+            checkpoint_alisases[id] = self

+

+    def calculate_shorthash(self):

+        self.sha256 = hashes.sha256(self.filename, "checkpoint/" + self.name)

+        self.shorthash = self.sha256[0:10]

+

+        if self.shorthash not in self.ids:

+            self.ids += [self.shorthash, self.sha256]

+            self.register()

+

+        self.title = f'{self.name} [{self.shorthash}]'

+

+        return self.shorthash

+

+

 try:

     # this silences the annoying "Some weights of the model checkpoint were not used when initializing..." message at start.

 

@@ -42,61 +87,50 @@ def setup_model():
     enable_midas_autodownload()

 

 

-def checkpoint_tiles(): 

-    convert = lambda name: int(name) if name.isdigit() else name.lower() 

-    alphanumeric_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)] 

-    return sorted([x.title for x in checkpoints_list.values()], key = alphanumeric_key)

-

-

-def list_models():

-    checkpoints_list.clear()

-    model_list = modelloader.load_models(model_path=model_path, command_path=shared.cmd_opts.ckpt_dir, ext_filter=[".ckpt", ".safetensors"])

-

-    def modeltitle(path, shorthash):

-        abspath = os.path.abspath(path)

+def checkpoint_tiles():

+    def convert(name):

+        return int(name) if name.isdigit() else name.lower()

 

-        if shared.cmd_opts.ckpt_dir is not None and abspath.startswith(shared.cmd_opts.ckpt_dir):

-            name = abspath.replace(shared.cmd_opts.ckpt_dir, '')

-        elif abspath.startswith(model_path):

-            name = abspath.replace(model_path, '')

-        else:

-            name = os.path.basename(path)

+    def alphanumeric_key(key):

+        return [convert(c) for c in re.split('([0-9]+)', key)]

 

-        if name.startswith("\\") or name.startswith("/"):

-            name = name[1:]

+    return sorted([x.title for x in checkpoints_list.values()], key=alphanumeric_key)

 

-        shortname = os.path.splitext(name.replace("/", "_").replace("\\", "_"))[0]

 

-        return f'{name} [{shorthash}]', shortname

+def list_models():

+    checkpoints_list.clear()

+    checkpoint_alisases.clear()

+    model_list = modelloader.load_models(model_path=model_path, command_path=shared.cmd_opts.ckpt_dir, ext_filter=[".ckpt", ".safetensors"], ext_blacklist=[".vae.safetensors"])

 

     cmd_ckpt = shared.cmd_opts.ckpt

     if os.path.exists(cmd_ckpt):

-        h = model_hash(cmd_ckpt)

-        title, short_model_name = modeltitle(cmd_ckpt, h)

-        checkpoints_list[title] = CheckpointInfo(cmd_ckpt, title, h, short_model_name, shared.cmd_opts.config)

-        shared.opts.data['sd_model_checkpoint'] = title

+        checkpoint_info = CheckpointInfo(cmd_ckpt)

+        checkpoint_info.register()

+

+        shared.opts.data['sd_model_checkpoint'] = checkpoint_info.title

     elif cmd_ckpt is not None and cmd_ckpt != shared.default_sd_model_file:

         print(f"Checkpoint in --ckpt argument not found (Possible it was moved to {model_path}: {cmd_ckpt}", file=sys.stderr)

+

     for filename in model_list:

-        h = model_hash(filename)

-        title, short_model_name = modeltitle(filename, h)

+        checkpoint_info = CheckpointInfo(filename)

+        checkpoint_info.register()

 

-        basename, _ = os.path.splitext(filename)

-        config = basename + ".yaml"

-        if not os.path.exists(config):

-            config = shared.cmd_opts.config

 

-        checkpoints_list[title] = CheckpointInfo(filename, title, h, short_model_name, config)

+def get_closet_checkpoint_match(search_string):

+    checkpoint_info = checkpoint_alisases.get(search_string, None)

+    if checkpoint_info is not None:

+        return checkpoint_info

 

+    found = sorted([info for info in checkpoints_list.values() if search_string in info.title], key=lambda x: len(x.title))

+    if found:

+        return found[0]

 

-def get_closet_checkpoint_match(searchString):

-    applicable = sorted([info for info in checkpoints_list.values() if searchString in info.title], key = lambda x:len(x.title))

-    if len(applicable) > 0:

-        return applicable[0]

     return None

 

 

 def model_hash(filename):

+    """old hash that only looks at a small part of the file and is prone to collisions"""

+

     try:

         with open(filename, "rb") as file:

             import hashlib

@@ -112,7 +146,7 @@ def model_hash(filename):
 def select_checkpoint():

     model_checkpoint = shared.opts.sd_model_checkpoint

         

-    checkpoint_info = checkpoints_list.get(model_checkpoint, None)

+    checkpoint_info = checkpoint_alisases.get(model_checkpoint, None)

     if checkpoint_info is not None:

         return checkpoint_info

 

@@ -168,7 +202,8 @@ def get_state_dict_from_checkpoint(pl_sd):
 def read_state_dict(checkpoint_file, print_global_state=False, map_location=None):

     _, extension = os.path.splitext(checkpoint_file)

     if extension.lower() == ".safetensors":

-        pl_sd = safetensors.torch.load_file(checkpoint_file, device=map_location or shared.weight_load_location)

+        device = map_location or shared.weight_load_location or devices.get_optimal_device_name()

+        pl_sd = safetensors.torch.load_file(checkpoint_file, device=device)

     else:

         pl_sd = torch.load(checkpoint_file, map_location=map_location or shared.weight_load_location)

 

@@ -179,59 +214,87 @@ def read_state_dict(checkpoint_file, print_global_state=False, map_location=None
     return sd

 

 

-def load_model_weights(model, checkpoint_info, vae_file="auto"):

-    checkpoint_file = checkpoint_info.filename

-    sd_model_hash = checkpoint_info.hash

-

-    cache_enabled = shared.opts.sd_checkpoint_cache > 0

+def get_checkpoint_state_dict(checkpoint_info: CheckpointInfo, timer):

+    sd_model_hash = checkpoint_info.calculate_shorthash()

+    timer.record("calculate hash")

 

-    if cache_enabled and checkpoint_info in checkpoints_loaded:

+    if checkpoint_info in checkpoints_loaded:

         # use checkpoint cache

         print(f"Loading weights [{sd_model_hash}] from cache")

-        model.load_state_dict(checkpoints_loaded[checkpoint_info])

-    else:

-        # load from file

-        print(f"Loading weights [{sd_model_hash}] from {checkpoint_file}")

+        return checkpoints_loaded[checkpoint_info]

 

-        sd = read_state_dict(checkpoint_file)

-        model.load_state_dict(sd, strict=False)

-        del sd

-        

-        if cache_enabled:

-            # cache newly loaded model

-            checkpoints_loaded[checkpoint_info] = model.state_dict().copy()

+    print(f"Loading weights [{sd_model_hash}] from {checkpoint_info.filename}")

+    res = read_state_dict(checkpoint_info.filename)

+    timer.record("load weights from disk")

+

+    return res

+

+

+def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer):

+    title = checkpoint_info.title

+    sd_model_hash = checkpoint_info.calculate_shorthash()

+    timer.record("calculate hash")

+

+    if checkpoint_info.title != title:

+        shared.opts.data["sd_model_checkpoint"] = checkpoint_info.title

+

+    if state_dict is None:

+        state_dict = get_checkpoint_state_dict(checkpoint_info, timer)

+

+    model.load_state_dict(state_dict, strict=False)

+    del state_dict

+    timer.record("apply weights to model")

+

+    if shared.opts.sd_checkpoint_cache > 0:

+        # cache newly loaded model

+        checkpoints_loaded[checkpoint_info] = model.state_dict().copy()

 

-        if shared.cmd_opts.opt_channelslast:

-            model.to(memory_format=torch.channels_last)

+    if shared.cmd_opts.opt_channelslast:

+        model.to(memory_format=torch.channels_last)

+        timer.record("apply channels_last")

 

-        if not shared.cmd_opts.no_half:

-            vae = model.first_stage_model

+    if not shared.cmd_opts.no_half:

+        vae = model.first_stage_model

+        depth_model = getattr(model, 'depth_model', None)

 

-            # with --no-half-vae, remove VAE from model when doing half() to prevent its weights from being converted to float16

-            if shared.cmd_opts.no_half_vae:

-                model.first_stage_model = None

+        # with --no-half-vae, remove VAE from model when doing half() to prevent its weights from being converted to float16

+        if shared.cmd_opts.no_half_vae:

+            model.first_stage_model = None

+        # with --upcast-sampling, don't convert the depth model weights to float16

+        if shared.cmd_opts.upcast_sampling and depth_model:

+            model.depth_model = None

 

-            model.half()

-            model.first_stage_model = vae

+        model.half()

+        model.first_stage_model = vae

+        if depth_model:

+            model.depth_model = depth_model

 

-        devices.dtype = torch.float32 if shared.cmd_opts.no_half else torch.float16

-        devices.dtype_vae = torch.float32 if shared.cmd_opts.no_half or shared.cmd_opts.no_half_vae else torch.float16

+        timer.record("apply half()")

 

-        model.first_stage_model.to(devices.dtype_vae)

+    devices.dtype = torch.float32 if shared.cmd_opts.no_half else torch.float16

+    devices.dtype_vae = torch.float32 if shared.cmd_opts.no_half or shared.cmd_opts.no_half_vae else torch.float16

+    devices.dtype_unet = model.model.diffusion_model.dtype

+    devices.unet_needs_upcast = shared.cmd_opts.upcast_sampling and devices.dtype == torch.float16 and devices.dtype_unet == torch.float16

+

+    model.first_stage_model.to(devices.dtype_vae)

+    timer.record("apply dtype to VAE")

 

     # clean up cache if limit is reached

-    if cache_enabled:

-        while len(checkpoints_loaded) > shared.opts.sd_checkpoint_cache + 1: # we need to count the current model

-            checkpoints_loaded.popitem(last=False)  # LRU

+    while len(checkpoints_loaded) > shared.opts.sd_checkpoint_cache:

+        checkpoints_loaded.popitem(last=False)

 

     model.sd_model_hash = sd_model_hash

-    model.sd_model_checkpoint = checkpoint_file

+    model.sd_model_checkpoint = checkpoint_info.filename

     model.sd_checkpoint_info = checkpoint_info

+    shared.opts.data["sd_checkpoint_hash"] = checkpoint_info.sha256

+

+    model.logvar = model.logvar.to(devices.device)  # fix for training

 

     sd_vae.delete_base_vae()

     sd_vae.clear_loaded_vae()

-    vae_file = sd_vae.resolve_vae(checkpoint_file, vae_file=vae_file)

-    sd_vae.load_vae(model, vae_file)

+    vae_file, vae_source = sd_vae.resolve_vae(checkpoint_info.filename)

+    sd_vae.load_vae(model, vae_file, vae_source)

+    timer.record("load VAE")

 

 

 def enable_midas_autodownload():

@@ -244,7 +307,7 @@ def enable_midas_autodownload():
     location automatically.

     """

 

-    midas_path = os.path.join(models_path, 'midas')

+    midas_path = os.path.join(paths.models_path, 'midas')

 

     # stable-diffusion-stability-ai hard-codes the midas model path to

     # a location that differs from where other scripts using this model look.

@@ -276,88 +339,144 @@ def enable_midas_autodownload():
 

     midas.api.load_model = load_model_wrapper

 

-def load_model(checkpoint_info=None):

+

+def repair_config(sd_config):

+

+    if not hasattr(sd_config.model.params, "use_ema"):

+        sd_config.model.params.use_ema = False

+

+    if shared.cmd_opts.no_half:

+        sd_config.model.params.unet_config.params.use_fp16 = False

+    elif shared.cmd_opts.upcast_sampling:

+        sd_config.model.params.unet_config.params.use_fp16 = True

+

+

+def load_model(checkpoint_info=None, already_loaded_state_dict=None, time_taken_to_load_state_dict=None):

     from modules import lowvram, sd_hijack

     checkpoint_info = checkpoint_info or select_checkpoint()

 

-    if checkpoint_info.config != shared.cmd_opts.config:

-        print(f"Loading config from: {checkpoint_info.config}")

-

     if shared.sd_model:

         sd_hijack.model_hijack.undo_hijack(shared.sd_model)

         shared.sd_model = None

         gc.collect()

         devices.torch_gc()

 

-    sd_config = OmegaConf.load(checkpoint_info.config)

-    

-    if should_hijack_inpainting(checkpoint_info):

-        # Hardcoded config for now...

-        sd_config.model.target = "ldm.models.diffusion.ddpm.LatentInpaintDiffusion"

-        sd_config.model.params.conditioning_key = "hybrid"

-        sd_config.model.params.unet_config.params.in_channels = 9

-        sd_config.model.params.finetune_keys = None

+    do_inpainting_hijack()

 

-        # Create a "fake" config with a different name so that we know to unload it when switching models.

-        checkpoint_info = checkpoint_info._replace(config=checkpoint_info.config.replace(".yaml", "-inpainting.yaml"))

+    timer = Timer()

 

-    if not hasattr(sd_config.model.params, "use_ema"):

-        sd_config.model.params.use_ema = False

+    if already_loaded_state_dict is not None:

+        state_dict = already_loaded_state_dict

+    else:

+        state_dict = get_checkpoint_state_dict(checkpoint_info, timer)

 

-    do_inpainting_hijack()

+    checkpoint_config = sd_models_config.find_checkpoint_config(state_dict, checkpoint_info)

 

-    if shared.cmd_opts.no_half:

-        sd_config.model.params.unet_config.params.use_fp16 = False

+    timer.record("find config")

+

+    sd_config = OmegaConf.load(checkpoint_config)

+    repair_config(sd_config)

+

+    timer.record("load config")

+

+    print(f"Creating model from config: {checkpoint_config}")

+

+    sd_model = None

+    try:

+        with sd_disable_initialization.DisableInitialization():

+            sd_model = instantiate_from_config(sd_config.model)

+    except Exception as e:

+        pass

 

-    sd_model = instantiate_from_config(sd_config.model)

-    load_model_weights(sd_model, checkpoint_info)

+    if sd_model is None:

+        print('Failed to create model quickly; will retry using slow method.', file=sys.stderr)

+        sd_model = instantiate_from_config(sd_config.model)

+

+    sd_model.used_config = checkpoint_config

+

+    timer.record("create model")

+

+    load_model_weights(sd_model, checkpoint_info, state_dict, timer)

 

     if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:

         lowvram.setup_for_low_vram(sd_model, shared.cmd_opts.medvram)

     else:

         sd_model.to(shared.device)

 

+    timer.record("move model to device")

+

     sd_hijack.model_hijack.hijack(sd_model)

 

+    timer.record("hijack")

+

     sd_model.eval()

     shared.sd_model = sd_model

 

+    sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings(force_reload=True)  # Reload embeddings after model load as they may or may not fit the model

+

+    timer.record("load textual inversion embeddings")

+

     script_callbacks.model_loaded_callback(sd_model)

 

-    print("Model loaded.")

+    timer.record("scripts callbacks")

+

+    print(f"Model loaded in {timer.summary()}.")

+

     return sd_model

 

 

 def reload_model_weights(sd_model=None, info=None):

     from modules import lowvram, devices, sd_hijack

     checkpoint_info = info or select_checkpoint()

- 

+

     if not sd_model:

         sd_model = shared.sd_model

 

-    if sd_model.sd_model_checkpoint == checkpoint_info.filename:

-        return

+    if sd_model is None:  # previous model load failed

+        current_checkpoint_info = None

+    else:

+        current_checkpoint_info = sd_model.sd_checkpoint_info

+        if sd_model.sd_model_checkpoint == checkpoint_info.filename:

+            return

+

+        if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:

+            lowvram.send_everything_to_cpu()

+        else:

+            sd_model.to(devices.cpu)

+

+        sd_hijack.model_hijack.undo_hijack(sd_model)

+

+    timer = Timer()

+

+    state_dict = get_checkpoint_state_dict(checkpoint_info, timer)

+

+    checkpoint_config = sd_models_config.find_checkpoint_config(state_dict, checkpoint_info)

+

+    timer.record("find config")

 

-    if sd_model.sd_checkpoint_info.config != checkpoint_info.config or should_hijack_inpainting(checkpoint_info) != should_hijack_inpainting(sd_model.sd_checkpoint_info):

+    if sd_model is None or checkpoint_config != sd_model.used_config:

         del sd_model

         checkpoints_loaded.clear()

-        load_model(checkpoint_info)

+        load_model(checkpoint_info, already_loaded_state_dict=state_dict, time_taken_to_load_state_dict=timer.records["load weights from disk"])

         return shared.sd_model

 

-    if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:

-        lowvram.send_everything_to_cpu()

-    else:

-        sd_model.to(devices.cpu)

-

-    sd_hijack.model_hijack.undo_hijack(sd_model)

+    try:

+        load_model_weights(sd_model, checkpoint_info, state_dict, timer)

+    except Exception as e:

+        print("Failed to load checkpoint, restoring previous")

+        load_model_weights(sd_model, current_checkpoint_info, None, timer)

+        raise

+    finally:

+        sd_hijack.model_hijack.hijack(sd_model)

+        timer.record("hijack")

 

-    load_model_weights(sd_model, checkpoint_info)

+        script_callbacks.model_loaded_callback(sd_model)

+        timer.record("script callbacks")

 

-    sd_hijack.model_hijack.hijack(sd_model)

-    script_callbacks.model_loaded_callback(sd_model)

+        if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram:

+            sd_model.to(devices.device)

+            timer.record("move model to device")

 

-    if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram:

-        sd_model.to(devices.device)

+    print(f"Weights loaded in {timer.summary()}.")

 

-    print("Weights loaded.")

     return sd_model