optimize diffusers memory handling

CHANGELOG.md

@@ -5,6 +5,8 @@
 - general:
   - caching of extra network information to enable much faster create/refresh operations  
     thanks @midcoastal
+- diffusers:
+  - redo "move model to cpu" logic to be more reliable
 
 ## Update for 2023-08-17
 

installer.py

@@ -288,12 +288,9 @@ def check_python():
         log.debug(f'Git {git_version.replace("git version", "").strip()}')
 
 
-# Intel hasn't released a corresponding torchvision wheel along with torch and
-# ipex wheels, so we have to install official pytorch torchvision as a W/A.
-# However, the latest torchvision explicitly requires torch version == 2.0.1,
-# which is incompatible with the Intel torch version 2.0.0a0. This will cause
-# intel torch to be uninstalled when pip scans the dependencies of torchvision.
-# This function will check the torch version and force installing Intel torch
+# Intel hasn't released a corresponding torchvision wheel along with torch and ipex wheels, so we have to install official pytorch torchvision as a W/A.
+# However, the latest torchvision explicitly requires torch version == 2.0.1, which is incompatible with the Intel torch version 2.0.0a0. This will cause
+# intel torch to be uninstalled when pip scans the dependencies of torchvision. This function will check the torch version and force installing Intel torch
 # 2.0.0a0 to avoid the underlying dll version error.
 # TODO(Disty or Nuullll) remove this W/A when Intel releases torchvision wheel for windows.
 def fix_ipex_win_torch():
@@ -306,8 +303,8 @@ def fix_ipex_win_torch():
             log.warning(f'Incompatible torch version {installed_torch_ver} for ipex windows, reinstalling to {ipex_torch_ver}')
             torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.0.0a0 intel_extension_for_pytorch==2.0.110+gitba7f6c1 -f https://developer.intel.com/ipex-whl-stable-xpu')
             install(torch_command)
-            import torch
-            import intel_extension_for_pytorch as ipex
+            import torch # pylint: disable=unused-import
+            import intel_extension_for_pytorch as ipex # pylint: disable=unused-import
     except Exception as e:
         log.warning(e)
 
@@ -340,7 +337,6 @@ def check_torch():
         log.info('AMD ROCm toolkit detected')
         os.environ.setdefault('PYTORCH_HIP_ALLOC_CONF', 'garbage_collection_threshold:0.8,max_split_size_mb:512')
         os.environ.setdefault('TENSORFLOW_PACKAGE', 'tensorflow-rocm')
-
         try:
             command = subprocess.run('rocm_agent_enumerator', shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
             amd_gpus = command.stdout.decode(encoding="utf8", errors="ignore").split('\n')
@@ -350,31 +346,26 @@ def check_torch():
             log.debug(f'Run rocm_agent_enumerator failed: {e}')
             amd_gpus = []
 
-        # use the first available amd gpu by default
-        hip_visible_devices = []
+        hip_visible_devices = [] # use the first available amd gpu by default
         for idx, gpu in enumerate(amd_gpus):
             if gpu in ['gfx1100', 'gfx1101', 'gfx1102']:
                 hip_visible_devices.append((idx, gpu, 'navi3x'))
                 break
-            # experimental navi 2x support
-            if gpu in ['gfx1030', 'gfx1031', 'gfx1032', 'gfx1034']:
+            if gpu in ['gfx1030', 'gfx1031', 'gfx1032', 'gfx1034']: # experimental navi 2x support
                 hip_visible_devices.append((idx, gpu, 'navi2x'))
                 break
         if len(hip_visible_devices) > 0:
             idx, gpu, arch = hip_visible_devices[0]
             log.debug(f'ROCm agent used by default: idx={idx} gpu={gpu} arch={arch}')
-
             os.environ.setdefault('HIP_VISIBLE_DEVICES', str(idx))
             if arch == 'navi3x':
                 os.environ.setdefault('HSA_OVERRIDE_GFX_VERSION', '11.0.0')
-                # do not use tensorflow-rocm for navi 3x
-                if os.environ.get('TENSORFLOW_PACKAGE') == 'tensorflow-rocm':
+                if os.environ.get('TENSORFLOW_PACKAGE') == 'tensorflow-rocm': # do not use tensorflow-rocm for navi 3x
                     os.environ['TENSORFLOW_PACKAGE'] = 'tensorflow==2.13.0'
             elif arch == 'navi2x':
                 os.environ.setdefault('HSA_OVERRIDE_GFX_VERSION', '10.3.0')
             else:
                 log.debug(f'HSA_OVERRIDE_GFX_VERSION auto config is skipped for {gpu}')
-
         try:
             command = subprocess.run('hipconfig --version', shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
             major_ver, minor_ver, *_ = command.stdout.decode(encoding="utf8", errors="ignore").split('.')
@@ -383,13 +374,11 @@ def check_torch():
         except Exception as e:
             log.debug(f'Run hipconfig failed: {e}')
             rocm_ver = None
-
         if rocm_ver in ['5.5', '5.6']:
             # install torch nightly via torchvision to avoid wasting bandwidth when torchvision depends on torch from yesterday
             torch_command = os.environ.get('TORCH_COMMAND', f'torchvision --pre --index-url https://download.pytorch.org/whl/nightly/rocm{rocm_ver}')
         else:
             torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.0.1 torchvision==0.15.2 --index-url https://download.pytorch.org/whl/rocm5.4.2')
-
         xformers_package = os.environ.get('XFORMERS_PACKAGE', 'none')
     elif allow_ipex and (args.use_ipex or shutil.which('sycl-ls') is not None or shutil.which('sycl-ls.exe') is not None or os.environ.get('ONEAPI_ROOT') is not None or os.path.exists('/opt/intel/oneapi') or os.path.exists("C:/Program Files (x86)/Intel/oneAPI") or os.path.exists("C:/oneAPI")):
         args.use_ipex = True # pylint: disable=attribute-defined-outside-init

modules/modelloader.py

@@ -14,14 +14,12 @@ def walk(top, onerror:callable=None):
     # A near-exact copy of `os.path.walk()`, trimmed slightly. Probably not nessesary for most people's collections, but makes a difference on really large datasets.
     nondirs = []
     walk_dirs = []
-
     try:
         scandir_it = os.scandir(top)
     except OSError as error:
         if onerror is not None:
             onerror(error, top)
         return
-
     with scandir_it:
         while True:
             try:
@@ -49,6 +47,8 @@ def walk(top, onerror:callable=None):
                         onerror(error, entry.path)
     # Recurse into sub-directories
     for new_path in walk_dirs:
+        if os.path.basename(new_path).startswith('models--'):
+            continue
         yield from walk(new_path, onerror)
     # Yield after recursion if going bottom up
     yield top, nondirs
@@ -214,7 +214,7 @@ def directory_directories(dir:str, *, recursive:bool=True) -> dict[str,tuple[flo
             except Exception:
                 pass
             del modelloader_directories[_dir]
-        for _dir, _files in walk(dir, lambda e, path: shared.log.error(f"Filesystem Walk Error: {e.__class__.__name__}({e}) -> {path}")):
+        for _dir, _files in walk(dir, lambda e, path: shared.log.debug(f"FS walk error: {e} {path}")):
             try:
                 mtime = os.path.getmtime(_dir)
                 if _dir not in modelloader_directories or mtime != modelloader_directories[_dir][0]:

modules/processing_diffusers.py

@@ -51,6 +51,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
             unet_device = model.unet.device
             model.unet.to(devices.cpu)
             devices.torch_gc()
+        model.vae.to(devices.device)
         latents.to(model.vae.device)
         decoded = model.vae.decode(latents / model.vae.config.scaling_factor, return_dict=False)[0]
         if shared.opts.diffusers_move_unet and not model.has_accelerate:
@@ -117,13 +118,9 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
         negative_pooled = None
         prompts, negative_prompts, prompts_2, negative_prompts_2 = fix_prompts(prompts, negative_prompts, prompts_2, negative_prompts_2)
         if shared.opts.prompt_attention in {'Compel parser', 'Full parser'}:
-            prompt_embed, pooled, negative_embed, negative_pooled = prompt_parser_diffusers.compel_encode_prompts(model,
-                                                                                                                  prompts,
-                                                                                                                  negative_prompts,
-                                                                                                                  prompts_2,
-                                                                                                                  negative_prompts_2,
-                                                                                                                  is_refiner,
-                                                                                                                  kwargs.pop("clip_skip", None))
+            prompt_embed, pooled, negative_embed, negative_pooled = prompt_parser_diffusers.compel_encode_prompts(model, prompts, negative_prompts,
+                                                                                                                  prompts_2, negative_prompts_2,
+                                                                                                                  is_refiner, kwargs.pop("clip_skip", None))
         if 'prompt' in possible:
             if hasattr(model, 'text_encoder') and 'prompt_embeds' in possible and prompt_embed is not None:
                 args['prompt_embeds'] = prompt_embed
@@ -261,7 +258,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
                 for i in range(len(decoded)):
                     images.save_image(decoded[i], path=p.outpath_samples, basename="", seed=seeds[i], prompt=prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-refiner")
 
-        if (shared.opts.diffusers_move_base or shared.cmd_opts.medvram or shared.opts.diffusers_model_cpu_offload) and not (shared.cmd_opts.lowvram or shared.opts.diffusers_seq_cpu_offload):
+        if shared.opts.diffusers_move_base and not shared.sd_model.has_accelerate:
             shared.log.debug('Diffusers: Moving base model to CPU')
             shared.sd_model.to(devices.cpu)
             devices.torch_gc()

modules/sd_models.py

@@ -696,7 +696,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
 
         if (shared.opts.diffusers_model_cpu_offload or shared.cmd_opts.medvram) and (shared.opts.diffusers_seq_cpu_offload or shared.cmd_opts.lowvram):
             shared.log.warning(f'Diffusers {op}: Model CPU offload (--medvram) and Sequential CPU offload (--lowvram) are not compatible')
-            shared.log.debug(f'Diffusers {op}: disable model CPU offload and --medvram')
+            shared.log.debug(f'Diffusers {op}: disabling model CPU offload and --medvram')
             shared.opts.diffusers_model_cpu_offload=False
             shared.cmd_opts.medvram=False
 
@@ -706,11 +706,21 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
         if hasattr(sd_model, "enable_model_cpu_offload"):
             if (shared.cmd_opts.medvram and devices.backend != "directml") or shared.opts.diffusers_model_cpu_offload:
                 shared.log.debug(f'Diffusers {op}: enable model CPU offload')
+                if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner:
+                    shared.opts.diffusers_move_base = False
+                    shared.opts.diffusers_move_unet = False
+                    shared.opts.diffusers_move_refiner = False
+                    shared.log.warning(f'Disabling {op} "Move model to CPU" since "Model CPU offload" is enabled')
                 sd_model.enable_model_cpu_offload()
                 sd_model.has_accelerate = True
         if hasattr(sd_model, "enable_sequential_cpu_offload"):
             if shared.cmd_opts.lowvram or shared.opts.diffusers_seq_cpu_offload:
                 shared.log.debug(f'Diffusers {op}: enable sequential CPU offload')
+                if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner:
+                    shared.opts.diffusers_move_base = False
+                    shared.opts.diffusers_move_unet = False
+                    shared.opts.diffusers_move_refiner = False
+                    shared.log.warning(f'Disabling {op} "Move model to CPU" since "Sequential CPU offload" is enabled')
                 sd_model.enable_sequential_cpu_offload(device=devices.device)
                 sd_model.has_accelerate = True
         if hasattr(sd_model, "enable_vae_slicing"):
@@ -760,10 +770,11 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
                 else:
                     if not refiner_enough_vram and not (shared.opts.diffusers_move_base and shared.opts.diffusers_move_refiner):
                         shared.log.warning(f"Insufficient GPU memory, using system memory as fallback: free={free_vram} GB")
-                        shared.log.debug('Enabled moving base model to CPU')
-                        shared.log.debug('Enabled moving refiner model to CPU')
-                        shared.opts.diffusers_move_base=True
-                        shared.opts.diffusers_move_refiner=True
+                        if not shared.opts.shared.opts.diffusers_seq_cpu_offload and not shared.opts.diffusers_model_cpu_offload:
+                            shared.log.debug('Enabled moving base model to CPU')
+                            shared.log.debug('Enabled moving refiner model to CPU')
+                            shared.opts.diffusers_move_base=True
+                            shared.opts.diffusers_move_refiner=True
                     shared.log.debug('Moving base model to CPU')
                     model_data.sd_model.to(devices.cpu)
                     devices.torch_gc(force=True)

modules/sd_vae.py

@@ -124,6 +124,9 @@ def resolve_vae(checkpoint_file):
             return vae_dict[basename], 'in VAE dir'
     else:
         vae_from_options = vae_dict.get(shared.opts.sd_vae, None) # 5th
+        if vae_from_options is not None:
+            return vae_from_options, 'specified in settings'
+        vae_from_options = vae_dict.get(shared.opts.sd_vae + '.safetensors', None) # 6th
         if vae_from_options is not None:
             return vae_from_options, 'specified in settings'
         shared.log.warning(f"VAE not found: {shared.opts.sd_vae}")
@@ -188,10 +191,8 @@ def load_vae_diffusers(model_file, vae_file=None, vae_source="from unknown sourc
         pass
     else:
         diffusers_load_config['variant'] = shared.opts.diffusers_vae_load_variant
-
     if shared.opts.diffusers_vae_upcast != 'default':
         diffusers_load_config['force_upcast'] = True if shared.opts.diffusers_vae_upcast == 'true' else False
-
     shared.log.debug(f'Diffusers VAE load config: {diffusers_load_config}')
     try:
         import diffusers
@@ -251,8 +252,11 @@ def reload_vae_weights(sd_model=None, vae_file=unspecified):
         load_vae(sd_model, vae_file, vae_source)
         sd_hijack.model_hijack.hijack(sd_model)
         script_callbacks.model_loaded_callback(sd_model)
+        if vae_file is not None:
+            shared.log.info(f"VAE weights loaded: {vae_file}")
+    # else:
+    #    load_vae_diffusers(model_file, vae_file, vae_source)
 
     if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram and not sd_model.has_accelerate:
         sd_model.to(devices.device)
-    shared.log.info(f"VAE weights loaded: {vae_file}")
     return sd_model

modules/shared.py

@@ -268,7 +268,7 @@ def list_themes():
 
 
 def disable_extensions():
-    if opts.lora_disable:
+    if opts.lyco_patch_lora:
         if 'Lora' not in opts.disabled_extensions:
             opts.data['disabled_extensions'].append('Lora')
     else:
@@ -403,8 +403,8 @@ options_templates.update(options_section(('diffusers', "Diffusers Settings"), {
     "diffusers_move_refiner": OptionInfo(True, "Move refiner model to CPU when not in use"),
     "diffusers_extract_ema": OptionInfo(True, "Use model EMA weights when possible"),
     "diffusers_generator_device": OptionInfo("default", "Generator device", gr.Radio, lambda: {"choices": ["default", "cpu"]}),
-    "diffusers_seq_cpu_offload": OptionInfo(False, "Enable sequential CPU offload"),
-    "diffusers_model_cpu_offload": OptionInfo(False, "Enable model CPU offload"),
+    "diffusers_model_cpu_offload": OptionInfo(False, "Enable model CPU offload (--medvram)"),
+    "diffusers_seq_cpu_offload": OptionInfo(False, "Enable sequential CPU offload (--lowvram)"),
     "diffusers_vae_upcast": OptionInfo("default", "VAE upcasting", gr.Radio, lambda: {"choices": ['default', 'true', 'false']}),
     "diffusers_vae_slicing": OptionInfo(True, "Enable VAE slicing"),
     "diffusers_vae_tiling": OptionInfo(False, "Enable VAE tiling"),
@@ -627,7 +627,7 @@ options_templates.update(options_section(('extra_networks', "Extra Networks"), {
     "extra_networks_card_fit": OptionInfo("cover", "UI image contain method", gr.Radio, lambda: {"choices": ["contain", "cover", "fill"]}),
     "extra_network_skip_indexing": OptionInfo(False, "Do not automatically build extra network pages", gr.Checkbox),
     "lyco_patch_lora": OptionInfo(False, "Use LyCoris handler for all Lora types", gr.Checkbox),
-    "lora_disable": OptionInfo(False, "Disable built-in Lora handler", gr.Checkbox, { "visible": True }, onchange=disable_extensions),
+    # "lora_disable": OptionInfo(False, "Disable built-in Lora handler", gr.Checkbox, { "visible": True }, onchange=disable_extensions),
     "lora_functional": OptionInfo(False, "Use Kohya method for handling multiple Loras", gr.Checkbox),
     "extra_networks_add_text_separator": OptionInfo(" ", "Extra text to add before <...> when adding extra network to prompt", gr.Text, { "visible": False }),
     "extra_networks_default_multiplier": OptionInfo(1.0, "Multiplier for extra networks", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),

modules/ui_extra_networks.py

@@ -11,9 +11,6 @@ from PIL import Image
 from modules import shared, scripts, modelloader
 from modules.generation_parameters_copypaste import image_from_url_text
 from modules.ui_components import ToolButton
-from logging import DEBUG
-from time import time
-from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn, SpinnerColumn
 
 extra_pages = []
 allowed_dirs = set()
@@ -161,8 +158,8 @@ class ExtraNetworksPage:
             return f"<div id='{tabname}_{self_name_id}_subdirs' class='extra-network-subdirs'></div><div id='{tabname}_{self_name_id}_cards' class='extra-network-cards'>Extra network page not ready<br>Click refresh to try again</div>"
         subdirs = {}
         allowed_folders = [os.path.abspath(x) for x in self.allowed_directories_for_previews()]
-        for parentdir, dirs in {dir: modelloader.directory_directories(dir) for dir in allowed_folders}.items():
-            for dir in dirs.keys():
+        for parentdir, dirs in {dir: modelloader.directory_directories(dir) for dir in allowed_folders}.items(): # pylint: disable=redefined-builtin
+            for dir in dirs.keys(): # pylint: disable=redefined-builtin
                 if shared.opts.diffusers_dir in dir:
                     subdirs[os.path.basename(shared.opts.diffusers_dir)] = 1
                 if 'models--' in dir:
@@ -191,23 +188,11 @@ class ExtraNetworksPage:
                 shared.log.error(f'Extra networks error listing items: {self.__class__}')
             self.create_xyz_grid()
             htmls = []
-            with Progress(
-                SpinnerColumn(),
-                TextColumn('[cyan]Creating Extra Network '+self.title+' HTML - {task.description}'),
-                BarColumn(), TaskProgressColumn(), TextColumn('({task.completed}/{task.total})'),
-                TimeRemainingColumn(), TimeElapsedColumn(), transient=not shared.log.isEnabledFor(DEBUG), expand=True
-            ) as progress:
-                task = progress.add_task(description='Initializing Items')
-                items = self.items
-                progress.update(task, total=len(items))
-                __t = time()
-                __i = 0
-                for item in items:
-                    __i += 1
-                    self.metadata[item["name"]] = item.get("metadata", {})
-                    self.info[item["name"]] = self.find_info(item['filename'])
-                    htmls.append(self.create_html_for_item(item, tabname))
-                    progress.update(task, advance=1, description=f"{round(__i/(shared.time.time()-__t))} item/s")
+            items = self.items
+            for item in items:
+                self.metadata[item["name"]] = item.get("metadata", {})
+                self.info[item["name"]] = self.find_info(item['filename'])
+                htmls.append(self.create_html_for_item(item, tabname))
             self.html += ''.join(htmls)
             if len(subdirs_html) > 0 or len(self.html) > 0:
                 res = f"<div id='{tabname}_{self_name_id}_subdirs' class='extra-network-subdirs'>{subdirs_html}</div><div id='{tabname}_{self_name_id}_cards' class='extra-network-cards'>{self.html}</div>"
@@ -262,7 +247,7 @@ class ExtraNetworksPage:
 
     def find_preview(self, path):
         preview_extensions = ["jpg", "jpeg", "png", "webp", "tiff", "jp2"]
-        dir = os.path.dirname(path)
+        dir = os.path.dirname(path) # pylint: disable=redefined-builtin
         paths = modelloader.directory_directories(dir, recursive=False)
         for file in [f'{path}.thumb.{ext}' for ext in preview_extensions]: # use thumbnail if exists
             if file in paths[dir][1] and os.path.exists(file):
@@ -274,7 +259,7 @@ class ExtraNetworksPage:
         return self.link_preview('html/card-no-preview.png')
 
     def find_description(self, path):
-        dir = os.path.dirname(path)
+        dir = os.path.dirname(path) # pylint: disable=redefined-builtin
         paths = modelloader.directory_directories(dir, recursive=False)
         for file in [f"{path}.txt", f"{path}.description.txt"]:
             if file in paths[dir][1]:
@@ -288,7 +273,7 @@ class ExtraNetworksPage:
         return None
 
     def find_info(self, path):
-        dir = os.path.dirname(path)
+        dir = os.path.dirname(path) # pylint: disable=redefined-builtin
         paths = modelloader.directory_directories(dir, recursive=False)
         basename, _ext = os.path.splitext(path)
         for file in [f"{path}.info", f"{path}.civitai.info", f"{basename}.info", f"{basename}.civitai.info"]:

modules/ui_extra_networks_checkpoints.py

@@ -16,6 +16,8 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
         checkpoint: sd_models.CheckpointInfo
         for name, checkpoint in sd_models.checkpoints_list.items():
             path, _ext = os.path.splitext(checkpoint.filename)
+            if not os.path.exists(path) and sd_models.model_path not in path:
+                path = os.path.abspath(os.path.join(checkpoint.path, os.pardir, os.pardir))
             yield {
                 "name": checkpoint.name_for_extra,
                 "filename": path,

requirements.txt

@@ -46,7 +46,7 @@ requests==2.31.0
 tqdm==4.65.0
 accelerate==0.20.3
 opencv-python-headless==4.7.0.72
-diffusers==0.19.3
+diffusers==0.20.0
 einops==0.4.1
 gradio==3.32.0
 huggingface_hub==0.16.4