Aimdo 0.5.5 + Auto-detect and enable --fast-disk when the disk is fast (CORE-440) (#16333)

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),
+        (FAST_PIN_SUBSETS, False, False, False),
-        (subsets, True, False),
+        (subsets, False, False, True),
+        (FAST_PIN_SUBSETS, True, False, False),
+        (subsets, True, False, True),
     ]
     if evict_active:
         tiers.extend([
-            (subsets, False, True),
+            (FAST_PIN_SUBSETS, False, True, False),
-            (subsets, True, True),
+            (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 = size + max(comfy.memory_management.RAM_CACHE_HEADROOM / 2, 2048 * 1024 ** 2) - comfy.system_memory.virtual_memory_available()
-        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()
     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):
+    for subsets, current_prompt, active, registrations in registration_eviction_tiers(evict_active):
-        shortfall -= free_model_pins(shortfall, subsets, current_prompt, active, registrations=True)
+        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" ])
+            model.partially_unload_ram(1e30, subsets=[ "patches", "patches-loaded", "patches-fast" ])
-            for subset in ("patches", "patches-loaded"):
+            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()

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):
+    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)
+        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:

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:

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)
 

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

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())

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

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):

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

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