diff --git a/comfy/model_management.py b/comfy/model_management.py index e62fd3d74..63a1d1549 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -659,6 +659,7 @@ def mark_mmap_dirty(storage): PIN_SUBSETS = [ "weights", "patches" ] LOADED_PIN_SUBSETS = [ "weights-loaded", "patches-loaded" ] +FAST_PIN_SUBSETS = [ "weights-fast", "patches-fast" ] def models_for_pin_eviction(active, current_prompt=None): for loaded_model in current_loaded_models: @@ -679,32 +680,48 @@ def free_model_pins(size, subsets, current_prompt, active, registrations=False): freed = model.unregister_inactive_pins(size, subsets=subsets) else: freed = model.partially_unload_ram(size, subsets=subsets) + if freed > 0: + detail( + "Pin eviction: model=%s subsets=%s workflow=%s active=%s action=%s freed_mb=%.1f", + model.model.__class__.__name__, subsets, current_prompt, active, + "unregister" if registrations else "destroy", freed / (1024 ** 2), + ) freed_total += freed size -= freed return freed_total def pin_eviction_tiers(loaded, evict_active): tiers = [ + (FAST_PIN_SUBSETS, False, False), (PIN_SUBSETS, False, None), (LOADED_PIN_SUBSETS, False, None), + (FAST_PIN_SUBSETS, True, False), (LOADED_PIN_SUBSETS, True, None), ] if not loaded: tiers.append((PIN_SUBSETS, True, False)) if evict_active: - tiers.append((PIN_SUBSETS, True, True)) + tiers.extend([ + (FAST_PIN_SUBSETS, False, True), + (FAST_PIN_SUBSETS, True, True), + (PIN_SUBSETS, True, True), + ]) return tiers def registration_eviction_tiers(evict_active): subsets = PIN_SUBSETS + LOADED_PIN_SUBSETS tiers = [ - (subsets, False, False), - (subsets, True, False), + (FAST_PIN_SUBSETS, False, False, False), + (subsets, False, False, True), + (FAST_PIN_SUBSETS, True, False, False), + (subsets, True, False, True), ] if evict_active: tiers.extend([ - (subsets, False, True), - (subsets, True, True), + (FAST_PIN_SUBSETS, False, True, False), + (subsets, False, True, True), + (FAST_PIN_SUBSETS, True, True, False), + (subsets, True, True, True), ]) return tiers @@ -730,10 +747,7 @@ def should_free_pins_for_ram_pressure(shortfall): def ensure_pin_budget(size, evict_active=False, loaded=False): if args.high_ram: return True - if args.fast_disk: - shortfall = TOTAL_PINNED_MEMORY + size - MAX_PINNED_MEMORY - else: - shortfall = size + max(comfy.memory_management.RAM_CACHE_HEADROOM / 2, 2048 * 1024 ** 2) - comfy.system_memory.virtual_memory_available() + shortfall = size + max(comfy.memory_management.RAM_CACHE_HEADROOM / 2, 2048 * 1024 ** 2) - comfy.system_memory.virtual_memory_available() if shortfall <= 0: return True @@ -747,8 +761,8 @@ def free_registrations(shortfall, evict_active=True): return True shortfall += REGISTERABLE_PIN_HYSTERESIS - for subsets, current_prompt, active in registration_eviction_tiers(evict_active): - shortfall -= free_model_pins(shortfall, subsets, current_prompt, active, registrations=True) + for subsets, current_prompt, active, registrations in registration_eviction_tiers(evict_active): + shortfall -= free_model_pins(shortfall, subsets, current_prompt, active, registrations=registrations) return shortfall <= REGISTERABLE_PIN_HYSTERESIS def ensure_pin_registerable(size, evict_active=True): @@ -1456,7 +1470,7 @@ def reset_cast_buffers(): pin_state = model.model.dynamic_pins[model.load_device] if pin_state["active"]: - for subset in ("weights", "weights-loaded"): + for subset in ("weights", "weights-loaded", "weights-fast"): *_, buckets = pin_state[subset] for size, bucket in list(buckets.items()): bucket[:] = [ entry for entry in bucket if entry[-1] is not None ] @@ -1464,8 +1478,8 @@ def reset_cast_buffers(): del buckets[size] pin_state["active"] = False - model.partially_unload_ram(1e30, subsets=[ "patches", "patches-loaded" ]) - for subset in ("patches", "patches-loaded"): + model.partially_unload_ram(1e30, subsets=[ "patches", "patches-loaded", "patches-fast" ]) + for subset in ("patches", "patches-loaded", "patches-fast"): pin_state[subset] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, pinned_hostbuf_size(model.model_size())), [], [-1], [0], [0], {}) STREAM_CAST_BUFFERS.clear() diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 08e69fae1..9c7e590a7 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -338,9 +338,10 @@ class LazyCastingParamPiece(torch.nn.Parameter): class ModelPatcher: - def __init__(self, model, load_device, offload_device, size=0, weight_inplace_update=False): + def __init__(self, model, load_device, offload_device, size=0, weight_inplace_update=False, fast_disk=False): self.size = size self.model = model + self.fast_disk = bool(comfy.model_management.args.fast_disk or fast_disk) if not hasattr(self.model, 'device'): logging.debug("Model doesn't have a device attribute.") self.model.device = offload_device @@ -442,7 +443,7 @@ class ModelPatcher: if model_override is None: model_override = self.get_clone_model_override() - n = class_(model_override[0], self.load_device, self.offload_device, self.model_size(), weight_inplace_update=self.weight_inplace_update) + n = class_(model_override[0], self.load_device, self.offload_device, self.model_size(), weight_inplace_update=self.weight_inplace_update, fast_disk=self.fast_disk) n.patches = {} for k in self.patches: n.patches[k] = self.patches[k][:] @@ -1748,14 +1749,14 @@ class ModelPatcher: class ModelPatcherDynamic(ModelPatcher): - def __new__(cls, model=None, load_device=None, offload_device=None, size=0, weight_inplace_update=False): + def __new__(cls, model=None, load_device=None, offload_device=None, size=0, weight_inplace_update=False, fast_disk=False): if load_device is not None and comfy.model_management.is_device_cpu(load_device): #reroute to default MP for CPUs - return ModelPatcher(model, load_device, offload_device, size, weight_inplace_update) + return ModelPatcher(model, load_device, offload_device, size, weight_inplace_update, fast_disk) return super().__new__(cls) - def __init__(self, model, load_device, offload_device, size=0, weight_inplace_update=False): - super().__init__(model, load_device, offload_device, size, weight_inplace_update) + def __init__(self, model, load_device, offload_device, size=0, weight_inplace_update=False, fast_disk=False): + super().__init__(model, load_device, offload_device, size, weight_inplace_update, fast_disk) if not hasattr(self.model, "dynamic_vbars"): self.model.dynamic_vbars = {} if not hasattr(self.model, "dynamic_pins"): @@ -1782,6 +1783,9 @@ class ModelPatcherDynamic(ModelPatcher): "patches": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0], [0], {}), "weights-loaded": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0], [0], {}), "patches-loaded": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0], [0], {}), + "weights-fast": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0], [0], {}), + "patches-fast": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0], [0], {}), + "fast_disk": self.fast_disk, "hostbufs_initialized": False, "failed": False, "active": False, @@ -1878,6 +1882,8 @@ class ModelPatcherDynamic(ModelPatcher): pin_state["patches"] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, hostbuf_size), [], [-1], [0], [0], {}) pin_state["weights-loaded"] = (comfy_aimdo.host_buffer.HostBuffer(0, 64 * 1024 * 1024, hostbuf_size), [], [-1], [0], [0], {}) pin_state["patches-loaded"] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, hostbuf_size), [], [-1], [0], [0], {}) + pin_state["weights-fast"] = (comfy_aimdo.host_buffer.HostBuffer(0, 64 * 1024 * 1024, hostbuf_size), [], [-1], [0], [0], {}) + pin_state["patches-fast"] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, hostbuf_size), [], [-1], [0], [0], {}) pin_state["hostbufs_initialized"] = True pin_state["failed"] = False pin_state["active"] = True @@ -2063,11 +2069,11 @@ class ModelPatcherDynamic(ModelPatcher): def loaded_ram_size(self): pin_state = self.model.dynamic_pins[self.load_device] - return pin_state["weights"][0].size + pin_state["weights-loaded"][0].size + return pin_state["weights"][0].size + pin_state["weights-loaded"][0].size + pin_state["weights-fast"][0].size def pinned_memory_size(self): pin_state = self.model.dynamic_pins[self.load_device] - return pin_state["weights"][3][0] + pin_state["weights-loaded"][3][0] + return pin_state["weights"][3][0] + pin_state["weights-loaded"][3][0] + pin_state["weights-fast"][3][0] def unregister_inactive_pins(self, ram_to_unload, subsets=[ "weights-loaded", "patches-loaded", "weights", "patches" ]): freed = 0 @@ -2096,7 +2102,7 @@ class ModelPatcherDynamic(ModelPatcher): return freed return freed - def partially_unload_ram(self, ram_to_unload, subsets=[ "weights-loaded", "patches-loaded", "weights", "patches" ]): + def partially_unload_ram(self, ram_to_unload, subsets=[ "weights-fast", "patches-fast", "weights-loaded", "patches-loaded", "weights", "patches" ]): freed = 0 pin_state = self.model.dynamic_pins[self.load_device] for subset in subsets: diff --git a/comfy/model_prefetch.py b/comfy/model_prefetch.py index d5952c8e9..418faa51c 100644 --- a/comfy/model_prefetch.py +++ b/comfy/model_prefetch.py @@ -210,7 +210,7 @@ def prefetch_queue_pop(queue, device, module, dtype=None, core=None, enable_grap registerable_size += lowvram_fn.memory_required() offload_stream, fully_faulted = comfy.ops.cast_modules_with_vbar(comfy_modules, None, device, None, True, return_faulted=True) - if not comfy.model_management.args.fast_disk: + if not (comfy_modules and comfy_modules[0]._pin_state["fast_disk"]): comfy.model_management.ensure_pin_registerable(registerable_size) comfy.model_management.sync_stream(device, offload_stream) if fully_faulted and dtype is not None: diff --git a/comfy/ops.py b/comfy/ops.py index 3089454aa..f4ac760c9 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -184,9 +184,10 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin needs_cast = False xfer_source = [ s.weight, s.bias ] - subset = "weights" + fast_disk = s._pin_state["fast_disk"] + subset = "weights-fast" if fast_disk else "weights" pin = comfy.pinned_memory.get_pin(s, subset=subset) - if pin is None and not args.fast_disk: + if pin is None and not fast_disk: loaded_pin = comfy.pinned_memory.get_pin(s, subset="weights-loaded") if loaded_pin is not None or signature is not None: subset = "weights-loaded" @@ -227,7 +228,7 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin if pin is not None: cast_maybe_lowvram_patch([pin], dest, offload_stream) return - if signature is None or not args.fast_disk or args.high_ram: + if signature is None or not fast_disk or args.high_ram: comfy.pinned_memory.pin_memory(m, subset=subset, size=size) pin = comfy.pinned_memory.get_pin(m, subset=subset) cast_maybe_lowvram_patch(source, pin, offload_stream, xfer_dest2=dest) @@ -242,14 +243,14 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin lowvram_dest = get_cast_buffer(lowvram_size) lowvram_source.prepare(lowvram_dest, None, copy=False, commit=True) - subset = "patches" + subset = "patches-fast" if fast_disk else "patches" pin = comfy.pinned_memory.get_pin(lowvram_source, subset=subset) - if pin is None: + if pin is None and not fast_disk: loaded_pin = comfy.pinned_memory.get_pin(lowvram_source, subset="patches-loaded") if loaded_pin is not None: subset = "patches-loaded" pin = loaded_pin - elif signature is not None and not args.fast_disk: + elif signature is not None: subset = "patches-loaded" handle_pin(lowvram_source, pin, lowvram_source, lowvram_dest, subset=subset, size=lowvram_size) diff --git a/comfy/pinned_memory.py b/comfy/pinned_memory.py index e9a9a70e2..b3057c620 100644 --- a/comfy/pinned_memory.py +++ b/comfy/pinned_memory.py @@ -84,6 +84,7 @@ def pin_memory(module, subset="weights", size=None): size = comfy.memory_management.vram_aligned_size([ module.weight, module.bias ]) registerable_size = size loaded = subset.endswith("-loaded") + fast = subset.endswith("-fast") priority = module_pin.get("balancer_priority") if priority is None: @@ -93,7 +94,7 @@ def pin_memory(module, subset="weights", size=None): comfy.memory_management.extra_ram_release(comfy.memory_management.RAM_CACHE_HEADROOM) if (not comfy.model_management.ensure_pin_budget(size, loaded=loaded) or - not comfy.model_management.ensure_pin_registerable(registerable_size)): + not comfy.model_management.ensure_pin_registerable(registerable_size, evict_active=not fast)): return _steal_pin(module, stack, buckets, size, priority, subset) offset = hostbuf.size @@ -105,7 +106,7 @@ def pin_memory(module, subset="weights", size=None): pin.untyped_storage()._comfy_hostbuf = hostbuf if torch.cuda.cudart().cudaHostRegister(pin.data_ptr(), size, 1) != 0: comfy.model_management.discard_cuda_async_error() - comfy.model_management.free_registrations(size) + comfy.model_management.free_registrations(size, evict_active=not fast) if torch.cuda.cudart().cudaHostRegister(pin.data_ptr(), size, 1) != 0: comfy.model_management.discard_cuda_async_error() del pin diff --git a/comfy/sd.py b/comfy/sd.py index 914e69d04..640d59352 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -36,6 +36,7 @@ import os import comfy.utils import comfy.ops import comfy.model_prefetch +import comfy.storage from . import clip_vision from . import gligen @@ -267,7 +268,7 @@ class CLIP: self.tokenizer = tokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data) te_disable_dynamic = disable_dynamic or getattr(self.cond_stage_model, "disable_offload", False) ModelPatcher = comfy.model_patcher.ModelPatcher if te_disable_dynamic else comfy.model_patcher.CoreModelPatcher - self.patcher = ModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device) + self.patcher = ModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device, fast_disk=comfy.storage.state_dict_fast_disk(state_dict)) #Match torch.float32 hardcode upcast in TE implemention self.patcher.set_model_compute_dtype(torch.float32) self.patcher.hook_mode = comfy.hooks.EnumHookMode.MinVram @@ -487,6 +488,7 @@ class CLIP: class VAE: def __init__(self, sd=None, device=None, config=None, dtype=None, metadata=None): + fast_disk = comfy.storage.state_dict_fast_disk(sd) is_seedvr2_vae = "decoder.up_blocks.2.upsamplers.0.upscale_conv.weight" in sd if not is_seedvr2_vae and 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format sd = diffusers_convert.convert_vae_state_dict(sd) @@ -1090,7 +1092,7 @@ class VAE: mp = comfy.model_patcher.CoreModelPatcher if self.disable_offload: mp = comfy.model_patcher.ModelPatcher - self.patcher = mp(self.first_stage_model, load_device=self.device, offload_device=offload_device) + self.patcher = mp(self.first_stage_model, load_device=self.device, offload_device=offload_device, fast_disk=fast_disk) m, u = self.first_stage_model.load_state_dict(sd, strict=False, assign=self.patcher.is_dynamic()) if len(m) > 0: @@ -2222,7 +2224,7 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c model = model_config.get_model(sd, diffusion_model_prefix, device=inital_load_device) ModelPatcher = comfy.model_patcher.ModelPatcher if disable_dynamic else comfy.model_patcher.CoreModelPatcher offload_device = model_options.get("offload_device", model_management.unet_offload_device()) - model_patcher = ModelPatcher(model, load_device=load_device, offload_device=offload_device) + model_patcher = ModelPatcher(model, load_device=load_device, offload_device=offload_device, fast_disk=comfy.storage.state_dict_fast_disk(sd)) model.load_model_weights(sd, diffusion_model_prefix, assign=model_patcher.is_dynamic()) if output_vae: @@ -2362,7 +2364,7 @@ def load_diffusion_model_state_dict(sd, model_options={}, metadata=None, disable model = model_config.get_model(new_sd, "") ModelPatcher = comfy.model_patcher.ModelPatcher if disable_dynamic else comfy.model_patcher.CoreModelPatcher - model_patcher = ModelPatcher(model, load_device=load_device, offload_device=offload_device) + model_patcher = ModelPatcher(model, load_device=load_device, offload_device=offload_device, fast_disk=comfy.storage.state_dict_fast_disk(new_sd)) if not model_management.is_device_cpu(offload_device): model.to(offload_device) model.load_model_weights(new_sd, "", assign=model_patcher.is_dynamic()) diff --git a/comfy/storage.py b/comfy/storage.py new file mode 100644 index 000000000..0cc555758 --- /dev/null +++ b/comfy/storage.py @@ -0,0 +1,107 @@ +import functools +import logging +import os +import platform +import re + +import comfy_aimdo.storage + + +_NVME_NAMESPACE = re.compile(r"^(nvme\d+)n\d+$") + + +def _read(path): + try: + with open(path, encoding="utf-8") as f: + return f.read().strip() + except OSError: + return None + + +def _physical_block_devices(name): + partition = f"/sys/class/block/{name}/partition" + if os.path.exists(partition): + name = os.path.basename(os.path.dirname(os.path.realpath(f"/sys/class/block/{name}"))) + + slaves = f"/sys/class/block/{name}/slaves" + try: + children = os.listdir(slaves) + except OSError: + children = [] + if children: + devices = [] + for child in children: + devices.extend(_physical_block_devices(child)) + return devices + return [name] + + +def _fast_nvme(name): + match = _NVME_NAMESPACE.match(name) + if match is None: + return False + controller = match.group(1) + speed = _read(f"/sys/class/nvme/{controller}/device/current_link_speed") + width = _read(f"/sys/class/nvme/{controller}/device/current_link_width") + if speed is None or width is None: + return None + try: + speed_gts = float(speed.split()[0]) + width = int(width) + except ValueError: + return None + return (speed_gts >= 8.0 and width >= 4) or (speed_gts >= 32.0 and width >= 2) + + +@functools.lru_cache(maxsize=None) +def _linux_fast_storage(device): + sys_device = f"/sys/dev/block/{os.major(device)}:{os.minor(device)}" + if not os.path.exists(sys_device): + return None + name = os.path.basename(os.path.realpath(sys_device)) + devices = _physical_block_devices(name) + results = [_fast_nvme(x) for x in devices] + if any(x is None for x in results): + return None + return all(results) + + +def fast_storage(path): + system = platform.system() + if system == "Linux": + try: + device = os.stat(os.path.realpath(path)).st_dev + except OSError: + return None + return _linux_fast_storage(device) + if system == "Windows": + return comfy_aimdo.storage.fast_disk(path) + return None + + +def annotate_state_dict(state_dict, path): + path = os.path.realpath(path) + for value in state_dict.values(): + untyped_storage = getattr(value, "untyped_storage", None) + if untyped_storage is not None: + untyped_storage()._comfy_source_path = path + + +def state_dict_fast_disk(state_dict): + state_dicts = state_dict if isinstance(state_dict, (list, tuple)) else (state_dict,) + paths = set() + for sd in state_dicts: + for value in sd.values(): + untyped_storage = getattr(value, "untyped_storage", None) + if untyped_storage is not None: + path = getattr(untyped_storage(), "_comfy_source_path", None) + if path is not None: + paths.add(path) + return model_fast_disk(sorted(paths)) if paths else False + + +def model_fast_disk(paths): + results = [fast_storage(path) for path in paths] + fast = bool(results) and all(result is True for result in results) + logging.info("Model storage policy: fast_disk=%s paths=%s", fast, [os.path.realpath(path) for path in paths]) + return fast diff --git a/comfy/utils.py b/comfy/utils.py index b31b70daa..573a9e455 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -23,6 +23,7 @@ import struct import ctypes import os import comfy.memory_management +import comfy.storage import safetensors.torch import numpy as np from PIL import Image @@ -200,6 +201,7 @@ def load_torch_file(ckpt, safe_load=False, device=None, return_metadata=False): sd = pl_sd else: sd = pl_sd + comfy.storage.annotate_state_dict(sd, ckpt) return (sd, metadata) if return_metadata else sd def save_torch_file(sd, ckpt, metadata=None): diff --git a/requirements.txt b/requirements.txt index 00e42041f..ae59fa053 100644 --- a/requirements.txt +++ b/requirements.txt @@ -23,7 +23,7 @@ SQLAlchemy>=2.0.0 filelock av>=17.0.0 comfy-kitchen==0.2.34 -comfy-aimdo==0.5.3 +comfy-aimdo==0.5.5 requests simpleeval>=1.0.0 blake3 diff --git a/tests-unit/comfy_test/storage_test.py b/tests-unit/comfy_test/storage_test.py new file mode 100644 index 000000000..6d2afc64c --- /dev/null +++ b/tests-unit/comfy_test/storage_test.py @@ -0,0 +1,26 @@ +from unittest import mock + +import comfy.storage + + +def test_fast_nvme_link_thresholds(): + cases = [ + ("8.0 GT/s PCIe", "4", True), + ("16.0 GT/s PCIe", "4", True), + ("32.0 GT/s PCIe", "2", True), + ("8.0 GT/s PCIe", "2", False), + ] + for speed, width, expected in cases: + with mock.patch.object(comfy.storage, "_read", side_effect=[speed, width]): + assert comfy.storage._fast_nvme("nvme0n1") is expected + + +def test_non_nvme_is_not_fast(): + assert comfy.storage._fast_nvme("sda") is False + + +def test_every_model_file_must_be_on_fast_storage(): + with mock.patch.object(comfy.storage, "fast_storage", side_effect=[True, False]): + assert comfy.storage.model_fast_disk(["first", "second"]) is False + with mock.patch.object(comfy.storage, "fast_storage", side_effect=[True, True]): + assert comfy.storage.model_fast_disk(["first", "second"]) is True