Auto2222 trim pass #2

Author: oxysoft

core/devicelib.py

@@ -89,6 +89,4 @@ def autocast(enable=True):
 
 # device = device
 
-device = get_optimal_device()
-
-print(device)
+device = get_optimal_device()

modules/stable_diffusion_auto1111/SDAttention.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+
+class SDAttention(Enum):
+    LDM = 0
+    SPLIT_BASUJINDAL = 1
+    SPLIT_INVOKE = 2
+    SPLIT_DOGGETT = 3
+    XFORMERS = 4

modules/stable_diffusion_auto1111/StableDiffusionPlugin.py

@@ -1,14 +1,9 @@
 # some of those options should not be changed at all because they would break the model, so I removed them from options.
-import math
 from collections import namedtuple
-from ctypes import Union
 from datetime import datetime
-from enum import Enum
 import platform
 
 import tqdm
-from einops import rearrange
-from numpy import einsum
 
 from core.jobs import Job, JobParams
 from SDJob import SDJob
@@ -18,6 +13,7 @@
 from SDCheckpointLoader import SDCheckpointLoader
 from core.printing import printerr
 from modules.stable_diffusion_auto1111.HypernetworkLoader import HypernetworkLoader
+from modules.stable_diffusion_auto1111.SDAttention import SDAttention
 from modules.stable_diffusion_auto1111.SDSampler import SDSampler
 from modules.stable_diffusion_auto1111.SDEmbeddingLoader import SDEmbeddingLoader
 from Hypernetwork import Hypernetwork
@@ -29,7 +25,7 @@
 from core.modellib import *
 from core.installing import git_clone, move_files
 from core.options import *
-from core.paths import repodir, modeldir
+from core.paths import repodir
 from core.plugins import Plugin
 
 from core import promptlib, devicelib, paths
@@ -40,7 +36,6 @@
 import ldm
 import ldm.modules.attention
 import ldm.modules.diffusionmodules.model
-from ldm.util import default
 
 ldm_crossattention_forward = ldm.modules.attention.CrossAttention.forward
 ldm_nonlinearity = ldm.modules.diffusionmodules.model.nonlinearity
@@ -51,13 +46,6 @@
 #     text = "<p>" + "<br>\n".join([f"{html.escape(x)}" for x in text.split('\n')]) + "</p>"
 #     return text
 
-class StableDiffusionAttention(Enum):
-    LDM = 0
-    SPLIT_BASUJINDAL = 1
-    SPLIT_INVOKE = 2
-    SPLIT_DOGGETT = 3
-    XFORMERS = 4
-
 
 class StableDiffusionPlugin(Plugin):
     class Options:
@@ -80,7 +68,7 @@ def __init__(self, plugin):
             self.no_half = True
             self.no_half_vae = False
             self.vae_override = ""
-            self.attention = StableDiffusionAttention.SPLIT_DOGGETT
+            self.attention = SDAttention.SPLIT_DOGGETT
             self.k_quantize = True
             self.always_batch_cond_uncond = True
 
@@ -121,8 +109,8 @@ def install(self):
         assert repo.is_dir() is not None, f"Couldn't find Stable Diffusion in {repo}"
 
         # TODO install xformers if enabled
-        if self.opt.attention == StableDiffusionAttention.XFORMERS:
-            import xformers.ops
+        if self.opt.attention == SDAttention.XFORMERS:
+            pass
 
         # TODO install mps if invoke attention
 
@@ -817,26 +805,26 @@ def set_ldm_overrides(self):
         if not invokeAI_mps_available and devicelib.device.type == 'mps':
             print("Cannot use InvokeAI cross attention optimization for MPS without psutil package, which is not installed.")
             print("Reverting to LDM.")
-            mode = StableDiffusionAttention.LDM
+            mode = SDAttention.LDM
 
-        if mode == StableDiffusionAttention.XFORMERS:
+        if mode == SDAttention.XFORMERS:
             if not (torch.version.cuda and (6, 0) <= torch.cuda.get_device_capability(devicelib.device) <= (8, 6)):
                 print("Cannot use xformers attention with the current CUDA version or GPU. Reverting to LDM")
-                mode = StableDiffusionAttention.LDM
+                mode = SDAttention.LDM
 
         # Apply the overrides
         # ----------------------------------------
-        if mode == StableDiffusionAttention.XFORMERS and torch.version.cuda and (6, 0) <= torch.cuda.get_device_capability(devicelib.device) <= (8, 6):
+        if mode == SDAttention.XFORMERS and torch.version.cuda and (6, 0) <= torch.cuda.get_device_capability(devicelib.device) <= (8, 6):
             print("Applying xformers cross attention optimization.")
             ldm.modules.attention.CrossAttention.forward = xformers_attention_forward
             ldm.modules.diffusionmodules.model.AttnBlock.forward = xformers_attnblock_forward
-        elif mode == StableDiffusionAttention.SPLIT_BASUJINDAL:
+        elif mode == SDAttention.SPLIT_BASUJINDAL:
             print("Applying cross attention optimization (Basujindal)")
             ldm.modules.attention.CrossAttention.forward = split_cross_attention_forward_basujindal(self)
-        elif mode == StableDiffusionAttention.SPLIT_INVOKE or not torch.cuda.is_available():
+        elif mode == SDAttention.SPLIT_INVOKE or not torch.cuda.is_available():
             print("Applying cross attention optimization (InvokeAI)")
             ldm.modules.attention.CrossAttention.forward = split_cross_attention_forward_invokeai(self)
-        elif mode == StableDiffusionAttention.SPLIT_DOGGETT:
+        elif mode == SDAttention.SPLIT_DOGGETT:
             print("Applying cross attention optimization (Doggettx)")
             ldm.modules.attention.CrossAttention.forward = split_cross_attention_forward_doggett(self)
             ldm.modules.diffusionmodules.model.AttnBlock.forward = cross_attention_attnblock_forward

modules/stable_diffusion_auto2222/StableDiffusionPlugin2.py

@@ -22,7 +22,7 @@ def f(*args, **kwargs):
 
 class StableDiffusionPlugin2(Plugin):
     def load(self):
-        modelloader.cleanup_models()
+        # modelloader.cleanup_models()
         sd_models.setup_model()
         # codeformer.setup_model(cmd_opts.codeformer_models_path)
         # gfpgan.setup_model(cmd_opts.gfpgan_models_path)

modules/stable_diffusion_auto2222/devices.py

@@ -76,7 +76,7 @@ def autocast(disable=False):
     if disable:
         return contextlib.nullcontext()
 
-    if dtype == torch.float32 or shared.cmd_opts.precision == "full":
+    if dtype == torch.float32 or shared.precision == "full":
         return contextlib.nullcontext()
 
     return torch.autocast("cuda")

modules/stable_diffusion_auto2222/interrogate.py

@@ -11,7 +11,7 @@
 from torchvision.transforms.functional import InterpolationMode
 
 import shared
-import devices, paths, lowvram
+import devices, lowvram
 
 blip_image_eval_size = 384
 blip_model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_caption_capfilt_large.pth'
@@ -68,14 +68,14 @@ def load_clip_model(self):
     def load(self):
         if self.blip_model is None:
             self.blip_model = self.load_blip_model()
-            if not shared.cmd_opts.no_half and not self.running_on_cpu:
+            if not shared.no_half and not self.running_on_cpu:
                 self.blip_model = self.blip_model.half()
 
         self.blip_model = self.blip_model.to(devices.device_interrogate)
 
         if self.clip_model is None:
             self.clip_model, self.clip_preprocess = self.load_clip_model()
-            if not shared.cmd_opts.no_half and not self.running_on_cpu:
+            if not shared.no_half and not self.running_on_cpu:
                 self.clip_model = self.clip_model.half()
 
         self.clip_model = self.clip_model.to(devices.device_interrogate)
@@ -134,7 +134,7 @@ def interrogate(self, pil_image):
 
         try:
 
-            if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
+            if shared.lowvram or shared.medvram:
                 lowvram.send_everything_to_cpu()
                 devices.torch_gc()
 
@@ -148,7 +148,7 @@ def interrogate(self, pil_image):
 
             clip_image = self.clip_preprocess(pil_image).unsqueeze(0).type(self.dtype).to(devices.device_interrogate)
 
-            precision_scope = torch.autocast if shared.cmd_opts.precision == "autocast" else contextlib.nullcontext
+            precision_scope = torch.autocast if shared.precision == "autocast" else contextlib.nullcontext
             with torch.no_grad(), precision_scope("cuda"):
                 image_features = self.clip_model.encode_image(clip_image).type(self.dtype)
 

modules/stable_diffusion_auto2222/modelloader.py

@@ -7,16 +7,15 @@
 # from basicsr.utils.download_util import load_file_from_url
 import shared
 from core.modellib import load_file_from_url
-from paths import script_path, models_path
 
 
-def load_models(model_path: str, model_url: str = None, command_path: str = None, ext_filter=None, download_name=None) -> list:
+def load_models(model_dir: str, model_url: str = None, command_path: str = None, ext_filter=None, download_name=None) -> list:
     """
     A one-and done loader to try finding the desired models in specified directories.
 
     @param download_name: Specify to download from model_url immediately.
     @param model_url: If no other models are found, this will be downloaded on upscale.
-    @param model_path: The location to store/find models in.
+    @param model_dir: The location to store/find models in.
     @param command_path: A command-line argument to search for models in first.
     @param ext_filter: An optional list of filename extensions to filter by
     @return: A list of paths containing the desired model(s)
@@ -29,15 +28,15 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None
     try:
         places = []
 
-        if command_path is not None and command_path != model_path:
+        if command_path is not None and command_path != model_dir:
             pretrained_path = os.path.join(command_path, 'experiments/pretrained_models')
             if os.path.exists(pretrained_path):
                 print(f"Appending path: {pretrained_path}")
                 places.append(pretrained_path)
             elif os.path.exists(command_path):
                 places.append(command_path)
 
-        places.append(model_path)
+        places.append(model_dir)
 
         for place in places:
             if os.path.exists(place):
@@ -54,7 +53,7 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None
 
         if model_url is not None and len(output) == 0:
             if download_name is not None:
-                dl = load_file_from_url(model_url, model_path, True, download_name)
+                dl = load_file_from_url(model_url, model_dir, True, download_name)
                 output.append(dl)
             else:
                 output.append(model_url)
@@ -74,26 +73,26 @@ def friendly_name(file: str):
     return model_name
 
 
-def cleanup_models():
-    # This code could probably be more efficient if we used a tuple list or something to store the src/destinations
-    # and then enumerate that, but this works for now. In the future, it'd be nice to just have every "model" scaler
-    # somehow auto-register and just do these things...
-    root_path = script_path
-    src_path = models_path
-    dest_path = os.path.join(models_path, "Stable-diffusion")
-    move_files(src_path, dest_path, ".ckpt")
-    src_path = os.path.join(root_path, "ESRGAN")
-    dest_path = os.path.join(models_path, "ESRGAN")
-    move_files(src_path, dest_path)
-    src_path = os.path.join(root_path, "gfpgan")
-    dest_path = os.path.join(models_path, "GFPGAN")
-    move_files(src_path, dest_path)
-    src_path = os.path.join(root_path, "SwinIR")
-    dest_path = os.path.join(models_path, "SwinIR")
-    move_files(src_path, dest_path)
-    src_path = os.path.join(root_path, "repositories/latent-diffusion/experiments/pretrained_models/")
-    dest_path = os.path.join(models_path, "LDSR")
-    move_files(src_path, dest_path)
+# def cleanup_models():
+#     # This code could probably be more efficient if we used a tuple list or something to store the src/destinations
+#     # and then enumerate that, but this works for now. In the future, it'd be nice to just have every "model" scaler
+#     # somehow auto-register and just do these things...
+#     root_path = script_path
+#     src_path = models_path
+#     dest_path = os.path.join(models_path, "Stable-diffusion")
+#     move_files(src_path, dest_path, ".ckpt")
+#     src_path = os.path.join(root_path, "ESRGAN")
+#     dest_path = os.path.join(models_path, "ESRGAN")
+#     move_files(src_path, dest_path)
+#     src_path = os.path.join(root_path, "gfpgan")
+#     dest_path = os.path.join(models_path, "GFPGAN")
+#     move_files(src_path, dest_path)
+#     src_path = os.path.join(root_path, "SwinIR")
+#     dest_path = os.path.join(models_path, "SwinIR")
+#     move_files(src_path, dest_path)
+#     src_path = os.path.join(root_path, "repositories/latent-diffusion/experiments/pretrained_models/")
+#     dest_path = os.path.join(models_path, "LDSR")
+#     move_files(src_path, dest_path)
 
 
 def move_files(src_path: str, dest_path: str, ext_filter: str = None):

modules/stable_diffusion_auto2222/paths.py

@@ -313,7 +313,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
     else:
         p.all_subseeds = [int(subseed) + x for x in range(len(p.all_prompts))]
 
-    if os.path.exists(cmd_opts.embeddings_dir) and not p.do_not_reload_embeddings:
+    if os.path.exists(shared.embeddings_dir) and not p.do_not_reload_embeddings:
         modules.stable_diffusion_auto2222.sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings()
 
     # if p.scripts is not None:
@@ -361,7 +361,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
 
             del samples_ddim
 
-            if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
+            if shared.lowvram or shared.medvram:
                 lowvram.send_everything_to_cpu()
 
             devices.torch_gc()