from __future__ import annotations import contextlib import functools import json import logging import os import struct from typing import Any, Callable import torch from safetensors import safe_open from .cleanup import ( _device_matches as _shared_device_matches, _should_force_cpu_offload, adaptive_headroom_bytes, trim_resident_vram, unload_loaded_model, ) from .residency import ( KIND_CHECKPOINT, KIND_CLIP, KIND_CLIP_VISION, KIND_CONTROLNET, KIND_MODEL, KIND_VAE, REGISTRY, ) _LOG = logging.getLogger(__name__) _PATCHED = False _ORIGINALS: dict[str, Callable[..., Any]] = {} _METADATA_CPU_KEY_SUFFIXES = ("spiece_model", "tekken_model", "comfy_quant") _UNET_PREFIX_CANDIDATES = ("model.diffusion_model.", "model.model.", "net.") _WARNED_PICKLE_GPU_PATHS: set[str] = set() _SAFE_TENSORS_COMPONENT_CACHE_MAX = 32 _STICKY_PROTECTION_VRAM_FLOOR_RATIO = 0.125 _STICKY_PROTECTION_VRAM_FLOOR_CEIL_BYTES = 16 * 1024 ** 3 _SAFETENSORS_DTYPE_MAP = { "BOOL": torch.bool, "U8": torch.uint8, "I8": torch.int8, "I16": torch.int16, "U16": getattr(torch, "uint16", torch.int32), "I32": torch.int32, "U32": getattr(torch, "uint32", torch.int64), "I64": torch.int64, "U64": getattr(torch, "uint64", torch.int64), "F16": torch.float16, "BF16": torch.bfloat16, "F32": torch.float32, "F64": torch.float64, "F8_E4M3FN": getattr(torch, "float8_e4m3fn", torch.float16), "F8_E5M2": getattr(torch, "float8_e5m2", torch.float16), } def _safetensors_header_cache_key(path: str) -> tuple[str, int, int]: stat = os.stat(path) return (os.path.abspath(path), stat.st_mtime_ns, stat.st_size) def _read_safetensors_header(path: str) -> tuple[dict[str, dict[str, Any]], dict[str, str] | None]: with open(path, "rb") as handle: header_size = struct.unpack(" torch.dtype: if code is None: return torch.float32 return _SAFETENSORS_DTYPE_MAP.get(code, torch.float32) def _tensor_nbytes_from_header( tensor_info: dict[str, Any], *, dtype_override: torch.dtype | None = None, ) -> int: dtype = dtype_override if dtype_override is not None else _torch_dtype_from_safetensors_code(tensor_info.get("dtype")) numel = 1 for dim in tensor_info.get("shape", ()): numel *= int(dim) return int(numel) * int(torch.empty((), dtype=dtype).element_size()) def _build_meta_state_dict_from_header(tensor_headers: dict[str, dict[str, Any]]) -> dict[str, torch.Tensor]: meta_state_dict: dict[str, torch.Tensor] = {} for key, tensor_info in tensor_headers.items(): shape = tuple(int(dim) for dim in tensor_info.get("shape", ())) meta_state_dict[key] = torch.empty( shape, dtype=_torch_dtype_from_safetensors_code(tensor_info.get("dtype")), device="meta", ) return meta_state_dict @functools.lru_cache(maxsize=_SAFE_TENSORS_COMPONENT_CACHE_MAX) def _cached_component_key_maps(cache_key: tuple[str, int, int]) -> dict[str, Any]: import comfy.sd as comfy_sd path = cache_key[0] tensor_headers, metadata = _read_safetensors_header(path) all_keys = tuple(tensor_headers) unet_prefix = infer_unet_prefix_from_keys(all_keys) meta_state_dict = _build_meta_state_dict_from_header(tensor_headers) model_config = comfy_sd.model_detection.model_config_from_unet(meta_state_dict, unet_prefix, metadata=metadata) def select_prefixed(prefixes: tuple[str, ...] | list[str] | None) -> tuple[tuple[str, str], ...]: if not prefixes: return () return tuple( (key, key) for key in all_keys if any(key.startswith(prefix) for prefix in prefixes) ) return { "metadata": metadata, "unet_prefix": unet_prefix, "model_config": model_config, "model": tuple((key, key[len(unet_prefix):]) for key in all_keys if key.startswith(unet_prefix)), "clip": select_prefixed(getattr(model_config, "text_encoder_key_prefix", None) or ()), "vae": select_prefixed(getattr(model_config, "vae_key_prefix", None) or ()), } def checkpoint_component_info_from_header(path: str) -> dict[str, Any] | None: try: return _cached_component_key_maps(_safetensors_header_cache_key(path)) except Exception as exc: _LOG.warning("GPU Resident Loader: failed to build selective safetensors header map for %s: %s", path, exc) return None def estimate_safetensors_tensor_bytes( path: str, *, selected_keys: list[str] | tuple[str, ...] | set[str] | None = None, dtype_override: torch.dtype | None = None, ) -> int | None: try: tensor_headers, _ = _read_safetensors_header(path) except Exception as exc: _LOG.warning("GPU Resident Loader: failed to estimate safetensors tensor bytes for %s: %s", path, exc) return None selected = None if selected_keys is None else set(selected_keys) total = 0 matched = 0 for key, tensor_info in tensor_headers.items(): if selected is not None and key not in selected: continue total += _tensor_nbytes_from_header(tensor_info, dtype_override=dtype_override) matched += 1 if selected is not None and selected and matched == 0: _LOG.warning( "GPU Resident Loader: selected tensor keys were provided but none matched %s; " "treating size as unknown for fallback", path, ) return None return int(total) def estimate_checkpoint_component_bytes( path: str, kind: str, *, dtype_override: torch.dtype | None = None, ) -> int | None: component_maps = checkpoint_component_info_from_header(path) if component_maps is None: return None pairs = component_maps.get(kind, ()) if not pairs: return None return estimate_safetensors_tensor_bytes( path, selected_keys=[source_key for source_key, _ in pairs], dtype_override=dtype_override, ) def _selected_component_keys_from_header(path: str, kind: str) -> dict[str, str] | None: if kind not in {KIND_MODEL, KIND_CLIP, KIND_VAE}: return None component_maps = checkpoint_component_info_from_header(path) if component_maps is None: return None pairs = component_maps.get(kind, ()) return dict(pairs) if pairs else None def _selected_component_suffix(kind: str | None) -> str | None: return { KIND_MODEL: "model_only", KIND_CLIP: "clip_only", KIND_VAE: "vae_only", }.get(kind) def _normalize_device(device: Any | None) -> torch.device | None: if device is None: return None if isinstance(device, torch.device): return device try: return torch.device(device) except Exception: return None def _device_string(device: torch.device | None) -> str: if device is None: return "auto" return str(device) def _devices_match(device_a: Any | None, device_b: Any | None) -> bool: return _shared_device_matches(device_a, device_b) def _cpu_offload_required(model: Any, loaded_device: Any | None) -> bool: current_device = None if hasattr(model, "current_loaded_device"): try: current_device = _normalize_device(model.current_loaded_device()) except Exception: current_device = None if current_device is None: current_device = _normalize_device(loaded_device) return _should_force_cpu_offload(model, active_device=current_device) @contextlib.contextmanager def _temporary_offload_device(model: Any, target_device: torch.device | None): if model is None or target_device is None or not hasattr(model, "offload_device"): yield False return original_device = getattr(model, "offload_device", None) if _devices_match(original_device, target_device): yield False return setattr(model, "offload_device", target_device) try: yield True finally: setattr(model, "offload_device", original_device) def _safe_open_device_arg(device: torch.device) -> Any: if device.type == "cuda": return device.index if device.index is not None else torch.cuda.current_device() if device.type == "cpu": return "cpu" return device.type def _copy_tensor_if_needed( tensor: torch.Tensor, target_device: torch.device, *, force_copy: bool = False, ) -> torch.Tensor: if tensor.device == target_device and not force_copy: return tensor return tensor.to(device=target_device, copy=True) def _tensor_key_requires_cpu(key: str) -> bool: return key.endswith(_METADATA_CPU_KEY_SUFFIXES) def _prepare_loaded_tensor( key: str, tensor: torch.Tensor, requested_device: torch.device, *, disable_mmap: bool, move_to_requested_device: bool = False, ) -> torch.Tensor: if _tensor_key_requires_cpu(key): return _copy_tensor_if_needed(tensor, torch.device("cpu")) if move_to_requested_device and tensor.device != requested_device: return _copy_tensor_if_needed(tensor, requested_device) if disable_mmap and tensor.device.type == "cpu": return _copy_tensor_if_needed(tensor, requested_device, force_copy=True) return tensor def _state_dict_device_summary(sd: dict[str, Any], requested_device: torch.device) -> str: devices: set[str] = set() for value in sd.values(): if torch.is_tensor(value): devices.add(str(value.device)) if not devices: return str(requested_device) if len(devices) == 1: return next(iter(devices)) return ", ".join(sorted(devices)) def _device_summary_from_observed(observed_devices: set[str], requested_device: torch.device) -> str: if not observed_devices: return str(requested_device) if len(observed_devices) == 1: return next(iter(observed_devices)) return ", ".join(sorted(observed_devices)) def infer_unet_prefix_from_keys(keys: list[str] | tuple[str, ...]) -> str: counts = {candidate: 0 for candidate in _UNET_PREFIX_CANDIDATES} for key in keys: for candidate in _UNET_PREFIX_CANDIDATES: if key.startswith(candidate): counts[candidate] += 1 break top = max(counts, key=counts.get) return top if counts[top] > 5 else "model." def load_safetensors_state_dict( ckpt: str, requested_device: torch.device, *, return_metadata: bool = False, selected_keys: dict[str, str] | None = None, ) -> tuple[dict[str, Any], Any, str, str]: import comfy.memory_management import comfy.utils as comfy_utils metadata = None if comfy.memory_management.aimdo_enabled and requested_device.type == "cpu" and selected_keys is None: sd, metadata = comfy_utils.load_safetensors(ckpt) if not return_metadata: metadata = None return sd, metadata, "cpu", "aimdo_cpu" disable_mmap = getattr(comfy_utils, "DISABLE_MMAP", False) def read_handle(device_arg: Any, *, move_to_requested_device: bool) -> tuple[dict[str, Any], Any, str]: observed_devices: set[str] = set() with safe_open(ckpt, framework="pt", device=device_arg) as handle: key_map = selected_keys if selected_keys is not None else {key: key for key in handle.keys()} sd: dict[str, Any] = {} for source_key, target_key in key_map.items(): tensor = handle.get_tensor(source_key) loaded = _prepare_loaded_tensor( source_key, tensor, requested_device, disable_mmap=disable_mmap, move_to_requested_device=move_to_requested_device, ) sd[target_key] = loaded if torch.is_tensor(loaded): observed_devices.add(str(loaded.device)) handle_metadata = handle.metadata() if return_metadata else None return sd, handle_metadata, _device_summary_from_observed(observed_devices, requested_device) try: safe_device = _safe_open_device_arg(requested_device) sd, metadata, actual_device = read_handle(safe_device, move_to_requested_device=False) return sd, metadata, actual_device, "direct" except Exception as exc: if requested_device.type == "cuda": _LOG.warning( "GPU Resident Loader: direct GPU safetensors load failed for %s; falling back to CPU path: %s", ckpt, exc, ) try: sd, metadata, actual_device = read_handle("cpu", move_to_requested_device=True) return sd, metadata, actual_device, "cpu_then_copy" except Exception as fallback_exc: raise fallback_exc from exc raise def _warn_pickle_gpu_compatibility(path: str, requested_device: torch.device) -> None: if requested_device.type != "cuda" or path in _WARNED_PICKLE_GPU_PATHS: return _WARNED_PICKLE_GPU_PATHS.add(path) _LOG.warning( "GPU Resident Loader: %s is not a safetensors file, so GPU-resident loading still goes through CPU-first torch.load(). " "Convert hot models with scripts/convert_checkpoint_to_safetensors.py for the narrow fast path.", path, ) def _sticky_protection_target(memory_required: int, device: Any) -> int: import comfy.model_management as model_management required = max(0, int(memory_required)) target = required + adaptive_headroom_bytes(required) minimum_inference_memory = getattr(model_management, "minimum_inference_memory", None) if callable(minimum_inference_memory): try: target = max(target, int(minimum_inference_memory())) except Exception: pass get_total_memory = getattr(model_management, "get_total_memory", None) if callable(get_total_memory): try: total_memory = int(get_total_memory(device)) except Exception: total_memory = 0 if total_memory > 0: target = max( target, min( _STICKY_PROTECTION_VRAM_FLOOR_CEIL_BYTES, int(total_memory * _STICKY_PROTECTION_VRAM_FLOOR_RATIO), ), ) return target def _resolved_context(kind: str, source_path: str | None) -> tuple[torch.device | None, str, str | None]: ctx = REGISTRY.current_context() if ctx is not None and ctx.explicit_device is not None: return ctx.explicit_device, ctx.kind, ctx.note return REGISTRY.explicit_load_device(kind=kind, source_path=source_path), kind, None def _record_generic_load( *, path: str, method: str, requested_device: torch.device | None, actual_device: str, note: str | None = None, error: str | None = None, ) -> None: ctx = REGISTRY.current_context() kind = ctx.kind if ctx is not None else "unknown" REGISTRY.record_load( path=path, kind=kind, method=method, requested_device=_device_string(requested_device), actual_device=actual_device, note=note or (ctx.note if ctx is not None else None), error=error, ) def _patched_load_torch_file(ckpt, safe_load=False, device=None, return_metadata=False): import comfy.memory_management import comfy.utils as comfy_utils requested_device = _normalize_device(device) ctx = REGISTRY.current_context() if requested_device is None and ctx is not None and ctx.explicit_device is not None: requested_device = ctx.explicit_device if requested_device is None: requested_device = torch.device("cpu") metadata = None lowered = str(ckpt).lower() if lowered.endswith((".safetensors", ".sft")): try: selected_keys = None selected_suffix = None if ctx is not None and ctx.kind in {KIND_MODEL, KIND_CLIP, KIND_VAE}: selected_keys = _selected_component_keys_from_header(ckpt, ctx.kind) selected_suffix = _selected_component_suffix(ctx.kind) sd, metadata, actual_device, load_mode = load_safetensors_state_dict( ckpt, requested_device, return_metadata=return_metadata, selected_keys=selected_keys, ) if load_mode == "aimdo_cpu": method = "safetensors_aimdo_cpu" elif selected_keys is not None and selected_suffix is not None: method = f"safetensors_cpu_then_copy_to_cuda_{selected_suffix}" if load_mode == "cpu_then_copy" else ( f"safetensors_gpu_direct_{selected_suffix}" if requested_device.type == "cuda" else f"safetensors_cpu_{selected_suffix}" ) else: method = "safetensors_cpu_then_copy_to_cuda" if load_mode == "cpu_then_copy" else ( "safetensors_gpu_direct" if requested_device.type == "cuda" else "safetensors_cpu" ) _record_generic_load( path=ckpt, method=method, requested_device=requested_device, actual_device=actual_device, ) return (sd, metadata) if return_metadata else sd except Exception as exc: if len(getattr(exc, "args", ())) > 0: message = exc.args[0] if isinstance(message, str): if "HeaderTooLarge" in message: raise ValueError( f"{message}\n\nFile path: {ckpt}\n\n" "The safetensors file is corrupt or invalid. Make sure this is actually a " "safetensors file and not a ckpt or pt or other filetype." ) from exc if "MetadataIncompleteBuffer" in message: raise ValueError( f"{message}\n\nFile path: {ckpt}\n\n" "The safetensors file is corrupt/incomplete. Check the file size and make sure " "you have copied/downloaded it correctly." ) from exc _record_generic_load( path=ckpt, method="safetensors_load_failed", requested_device=requested_device, actual_device="error", error=str(exc), ) raise torch_args = {} if getattr(comfy_utils, "MMAP_TORCH_FILES", False): torch_args["mmap"] = True _warn_pickle_gpu_compatibility(ckpt, requested_device) torch_load_device = torch.device("cpu") if requested_device.type == "cuda" else requested_device pl_sd = torch.load(ckpt, map_location=torch_load_device, weights_only=True, **torch_args) if "state_dict" in pl_sd: sd = pl_sd["state_dict"] else: if len(pl_sd) == 1: key = list(pl_sd.keys())[0] sd = pl_sd[key] if not isinstance(sd, dict): sd = pl_sd else: sd = pl_sd if isinstance(sd, dict): for key, value in list(sd.items()): if torch.is_tensor(value): sd[key] = _prepare_loaded_tensor( key, value, requested_device, disable_mmap=False, move_to_requested_device=True, ) method = "torch_load_cpu_first_to_cuda" if requested_device.type == "cuda" else "torch_load_cpu" _record_generic_load( path=ckpt, method=method, requested_device=requested_device, actual_device=_state_dict_device_summary(sd, requested_device) if isinstance(sd, dict) else str(requested_device), ) return (sd, metadata) if return_metadata else sd def _bind_checkpoint_outputs(result, source_path: str) -> None: if not result: return model = result[0] if len(result) > 0 else None clip = result[1] if len(result) > 1 else None vae = result[2] if len(result) > 2 else None if model is not None: REGISTRY.bind_object(model, source_path=source_path, kind=KIND_MODEL, note="checkpoint model") if clip is not None and getattr(clip, "patcher", None) is not None: REGISTRY.bind_object(clip.patcher, source_path=source_path, kind=KIND_CLIP, note="checkpoint clip") if vae is not None and getattr(vae, "patcher", None) is not None: REGISTRY.bind_object(vae.patcher, source_path=source_path, kind=KIND_VAE, note="checkpoint vae") def _bind_diffusers_outputs(result, source_path: str) -> None: if not result: return model = result[0] if len(result) > 0 else None clip = result[1] if len(result) > 1 else None vae = result[2] if len(result) > 2 else None if model is not None: REGISTRY.bind_object(model, source_path=source_path, kind=KIND_MODEL, note="diffusers model") if clip is not None and getattr(clip, "patcher", None) is not None: REGISTRY.bind_object(clip.patcher, source_path=source_path, kind=KIND_CLIP, note="diffusers clip") if vae is not None and getattr(vae, "patcher", None) is not None: REGISTRY.bind_object(vae.patcher, source_path=source_path, kind=KIND_VAE, note="diffusers vae") def _wrap_with_load_context(kind: str, path_arg_index: int = 0, bind_output: Callable[[Any, str], None] | None = None): def decorator(func: Callable[..., Any]) -> Callable[..., Any]: @functools.wraps(func) def wrapper(*args, **kwargs): source_path = None if len(args) > path_arg_index: source_path = args[path_arg_index] explicit_device = REGISTRY.explicit_load_device(kind=kind, source_path=source_path) with REGISTRY.load_context(kind=kind, source_path=source_path, explicit_device=explicit_device): result = func(*args, **kwargs) if bind_output is not None and source_path is not None: bind_output(result, source_path) return result return wrapper return decorator def _wrap_load_clip(func: Callable[..., Any]) -> Callable[..., Any]: @functools.wraps(func) def wrapper(*args, **kwargs): ckpt_paths = args[0] if args else kwargs.get("ckpt_paths") source_path = None if isinstance(ckpt_paths, (list, tuple)) and ckpt_paths: source_path = ckpt_paths[0] explicit_device = REGISTRY.explicit_load_device(kind=KIND_CLIP, source_path=source_path) with REGISTRY.load_context(kind=KIND_CLIP, source_path=source_path, explicit_device=explicit_device): clip = func(*args, **kwargs) if clip is not None and getattr(clip, "patcher", None) is not None and source_path is not None: REGISTRY.bind_object(clip.patcher, source_path=source_path, kind=KIND_CLIP) return clip return wrapper def _sticky_safe_batch_number(*, batch_count: int, free_memory: int, memory_used: int, device: Any) -> int: batches = max(1, int(max(0, int(free_memory)) / max(1, int(memory_used)))) batches = min(max(1, int(batch_count)), batches) if REGISTRY.get_policy() != "sticky_gpu" or device is None: return batches reserve = max(0, _sticky_protection_target(memory_used, device) - max(0, int(memory_used))) safe_budget = max(0, int(free_memory) - reserve) safe_batches = max(1, int(safe_budget / max(1, int(memory_used)))) capped = min(batches, safe_batches) if capped < batches: _LOG.debug( "GPU Resident Loader: capped VAE batch from %s to %s to preserve %s bytes of transient headroom.", batches, capped, reserve, ) return max(1, capped) def _wrap_vae_encode(func: Callable[..., Any]) -> Callable[..., Any]: @functools.wraps(func) def wrapper(self, pixel_samples): if REGISTRY.get_policy() != "sticky_gpu": return func(self, pixel_samples) import comfy.model_management as model_management self.throw_exception_if_invalid() pixel_samples = self.vae_encode_crop_pixels(pixel_samples) pixel_samples = pixel_samples.movedim(-1, 1) do_tile = False if self.latent_dim == 3 and pixel_samples.ndim < 5: if not self.not_video: pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0) else: pixel_samples = pixel_samples.unsqueeze(2) try: memory_used = self.memory_used_encode(pixel_samples.shape, self.vae_dtype) model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload) free_memory = self.patcher.get_free_memory(self.device) batch_number = _sticky_safe_batch_number( batch_count=pixel_samples.shape[0], free_memory=free_memory, memory_used=memory_used, device=self.device, ) samples = None for x in range(0, pixel_samples.shape[0], batch_number): pixels_in = self.process_input(pixel_samples[x:x + batch_number]).to(self.vae_dtype) if getattr(self.first_stage_model, "comfy_has_chunked_io", False): out = self.first_stage_model.encode(pixels_in, device=self.device) else: pixels_in = pixels_in.to(self.device) out = self.first_stage_model.encode(pixels_in) out = out.to(self.output_device).to(dtype=self.vae_output_dtype()) if samples is None: samples = torch.empty( (pixel_samples.shape[0],) + tuple(out.shape[1:]), device=self.output_device, dtype=self.vae_output_dtype(), ) samples[x:x + batch_number] = out except Exception as e: model_management.raise_non_oom(e) _LOG.warning("Warning: Ran out of memory when regular VAE encoding, retrying with tiled VAE encoding.") do_tile = True if do_tile: model_management.soft_empty_cache() if self.latent_dim == 3: tile = 256 overlap = tile // 4 samples = self.encode_tiled_3d(pixel_samples, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap)) elif self.latent_dim == 1 or self.extra_1d_channel is not None: samples = self.encode_tiled_1d(pixel_samples) else: samples = self.encode_tiled_(pixel_samples) return samples return wrapper def _wrap_vae_decode(func: Callable[..., Any]) -> Callable[..., Any]: @functools.wraps(func) def wrapper(self, samples_in, vae_options={}): if REGISTRY.get_policy() != "sticky_gpu": return func(self, samples_in, vae_options) import comfy.model_management as model_management self.throw_exception_if_invalid() pixel_samples = None do_tile = False if self.latent_dim == 2 and samples_in.ndim == 5: samples_in = samples_in[:, :, 0] try: memory_used = self.memory_used_decode(samples_in.shape, self.vae_dtype) model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload) free_memory = self.patcher.get_free_memory(self.device) batch_number = _sticky_safe_batch_number( batch_count=samples_in.shape[0], free_memory=free_memory, memory_used=memory_used, device=self.device, ) preallocated = False if getattr(self.first_stage_model, "comfy_has_chunked_io", False): pixel_samples = torch.empty( self.first_stage_model.decode_output_shape(samples_in.shape), device=self.output_device, dtype=self.vae_output_dtype(), ) preallocated = True for x in range(0, samples_in.shape[0], batch_number): samples = samples_in[x:x + batch_number].to(device=self.device, dtype=self.vae_dtype) if preallocated: self.first_stage_model.decode(samples, output_buffer=pixel_samples[x:x + batch_number], **vae_options) else: out = self.first_stage_model.decode(samples, **vae_options).to( device=self.output_device, dtype=self.vae_output_dtype(), copy=True, ) if pixel_samples is None: pixel_samples = torch.empty( (samples_in.shape[0],) + tuple(out.shape[1:]), device=self.output_device, dtype=self.vae_output_dtype(), ) pixel_samples[x:x + batch_number].copy_(out) del out self.process_output(pixel_samples[x:x + batch_number]) except Exception as e: model_management.raise_non_oom(e) _LOG.warning("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.") do_tile = True if do_tile: model_management.soft_empty_cache() dims = samples_in.ndim - 2 if dims == 1 or self.extra_1d_channel is not None: pixel_samples = self.decode_tiled_1d(samples_in) elif dims == 2: pixel_samples = self.decode_tiled_(samples_in) elif dims == 3: tile = 256 // self.spacial_compression_decode() overlap = tile // 4 pixel_samples = self.decode_tiled_3d(samples_in, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap)) pixel_samples = pixel_samples.to(self.output_device).movedim(1, -1) return pixel_samples return wrapper def _wrap_load_models_gpu(func: Callable[..., Any]) -> Callable[..., Any]: @functools.wraps(func) def wrapper(models, *args, **kwargs): import comfy.model_management as model_management requested_models = set() for model in list(models): requested_models.add(model) for additional in model.model_patches_models(): requested_models.add(additional) clone_conflicts: list[Any] = [] seen_loaded_ids: set[int] = set() for requested in requested_models: is_clone = getattr(requested, "is_clone", None) if not callable(is_clone): continue for loaded in list(model_management.current_loaded_models): try: dead = loaded.is_dead() except Exception: dead = False if id(loaded) in seen_loaded_ids or dead: continue loaded_model = getattr(loaded, "model", None) if loaded_model is None or loaded_model is requested: continue try: if not requested.is_clone(loaded_model): continue except Exception as exc: raise RuntimeError( "GPU Resident Loader: failed to evaluate clone-conflict state before replacement" ) from exc clone_conflicts.append(loaded) seen_loaded_ids.add(id(loaded)) clone_conflicts_unloaded = 0 try: for loaded in clone_conflicts: # ComfyUI's built-in clone replacement pops the wrapper and only calls detach(False), # which does not unpatch base weights. Fully unload before replacement or fail closed. if not unload_loaded_model( loaded, active_device=getattr(loaded, "device", None), force_offload_to_cpu=True, ): raise RuntimeError("GPU Resident Loader: failed to fully unload a clone-conflict wrapper before replacement") try: model_management.current_loaded_models.remove(loaded) except ValueError: pass clone_conflicts_unloaded += 1 finally: if clone_conflicts_unloaded > 0: if hasattr(model_management, "soft_empty_cache"): model_management.soft_empty_cache() REGISTRY.refresh_runtime_state() result = func(models, *args, **kwargs) for model in list(models): REGISTRY.touch(model) REGISTRY.refresh_runtime_state() return result return wrapper def _wrap_loaded_model_unload(func: Callable[..., Any]) -> Callable[..., Any]: @functools.wraps(func) def wrapper(self, memory_to_free=None, unpatch_weights=True): model = getattr(self, "model", None) loaded_device = getattr(self, "device", None) if not _cpu_offload_required(model, loaded_device): return func(self, memory_to_free=memory_to_free, unpatch_weights=unpatch_weights) with _temporary_offload_device(model, torch.device("cpu")) as redirected: if redirected: _LOG.debug( "GPU Resident Loader: redirecting unload of %s from %s to CPU to reclaim VRAM", type(getattr(model, "model", model)).__name__, _device_string(_normalize_device(loaded_device)), ) return func(self, memory_to_free=memory_to_free, unpatch_weights=unpatch_weights) return wrapper def _wrap_model_patcher_detach(func: Callable[..., Any]) -> Callable[..., Any]: @functools.wraps(func) def wrapper(self, unpatch_all=True): if not _cpu_offload_required(self, getattr(self.model, "device", None)): return func(self, unpatch_all=unpatch_all) with _temporary_offload_device(self, torch.device("cpu")) as redirected: if redirected: _LOG.debug( "GPU Resident Loader: redirecting detach of %s from %s to CPU to reclaim VRAM", type(getattr(self, "model", self)).__name__, _device_string(_normalize_device(getattr(self.model, "device", None))), ) return func(self, unpatch_all=unpatch_all) return wrapper def _wrap_free_memory(func: Callable[..., Any]) -> Callable[..., Any]: @functools.wraps(func) def wrapper(memory_required, device, keep_loaded=None, *args, **kwargs): import comfy.model_management as model_management keep_loaded = list(keep_loaded or []) protected_wrappers: list[Any] = [] sticky_wrappers: list[Any] = [] if REGISTRY.get_policy() == "sticky_gpu" and device is not None: sticky_wrappers = [w for w in REGISTRY.sticky_loaded_wrappers(device) if w not in keep_loaded] if device is not None and sticky_wrappers: protection_target = _sticky_protection_target(memory_required, device) keep_models = tuple( model for model in (getattr(loaded_wrapper, "model", None) for loaded_wrapper in keep_loaded) if model is not None ) try: trim_resident_vram( device=device, target_free_vram_bytes=protection_target, respect_sticky=True, sticky_floor_priority=0, allow_partial_unload=True, keep_models=keep_models, ) except Exception as exc: _LOG.debug("GPU Resident Loader: sticky pre-trim failed for free_memory(%s): %s", memory_required, exc) sticky_wrappers = [w for w in REGISTRY.sticky_loaded_wrappers(device) if w not in keep_loaded] try: free_now = model_management.get_free_memory(device) except Exception: free_now = None if free_now is None: protected_wrappers = sticky_wrappers else: unloadable_wrappers = [] for loaded in list(model_management.current_loaded_models): if loaded.device == device and loaded not in keep_loaded and not loaded.is_dead(): unloadable_wrappers.append(loaded) available_for_protection = max( 0, free_now + sum(max(0, loaded.model_loaded_memory()) for loaded in unloadable_wrappers) - protection_target, ) protected_memory = 0 for loaded in sticky_wrappers: estimated_memory = max(0, loaded.model_loaded_memory()) if protected_memory + estimated_memory <= available_for_protection: protected_wrappers.append(loaded) protected_memory += estimated_memory unloaded = func(memory_required, device, keep_loaded + protected_wrappers, *args, **kwargs) REGISTRY.refresh_runtime_state() return unloaded return wrapper def _remember_original(key: str, value: Callable[..., Any]) -> Callable[..., Any]: return _ORIGINALS.setdefault(key, value) def _patch_model_management_devices() -> None: import comfy.model_management as model_management kind_by_function = { "unet_offload_device": KIND_MODEL, "unet_inital_load_device": KIND_MODEL, "text_encoder_offload_device": KIND_CLIP, "text_encoder_device": KIND_CLIP, "vae_offload_device": KIND_VAE, "vae_device": KIND_VAE, } def wrap_device_func(name: str) -> None: key = f"model_management.{name}" original = _remember_original(key, getattr(model_management, name)) if getattr(model_management, name) is not original: return @functools.wraps(original) def wrapper(*args, **kwargs): result = original(*args, **kwargs) kind = kind_by_function.get(name) if not REGISTRY.wants_gpu_offload(kind): return result gpu_device = model_management.get_torch_device() if getattr(gpu_device, "type", None) == "cpu": return result return gpu_device setattr(model_management, name, wrapper) for name in ( "unet_offload_device", "text_encoder_offload_device", "vae_offload_device", "text_encoder_device", "vae_device", "unet_inital_load_device", ): if hasattr(model_management, name): wrap_device_func(name) original_free_memory = _remember_original("model_management.free_memory", model_management.free_memory) if model_management.free_memory is original_free_memory: model_management.free_memory = _wrap_free_memory(original_free_memory) original_load_models_gpu = _remember_original("model_management.load_models_gpu", model_management.load_models_gpu) if model_management.load_models_gpu is original_load_models_gpu: model_management.load_models_gpu = _wrap_load_models_gpu(original_load_models_gpu) original_model_unload = _remember_original("model_management.LoadedModel.model_unload", model_management.LoadedModel.model_unload) if model_management.LoadedModel.model_unload is original_model_unload: model_management.LoadedModel.model_unload = _wrap_loaded_model_unload(original_model_unload) def install_patches() -> None: global _PATCHED if _PATCHED: return import comfy.clip_vision as clip_vision import comfy.controlnet as controlnet import comfy.diffusers_load as diffusers_load import comfy.model_management as model_management import comfy.model_patcher as model_patcher import comfy.sd as comfy_sd import comfy.utils as comfy_utils original_load_torch_file = _remember_original("utils.load_torch_file", comfy_utils.load_torch_file) if comfy_utils.load_torch_file is original_load_torch_file: comfy_utils.load_torch_file = _patched_load_torch_file if hasattr(clip_vision, "load_torch_file"): original_clip_vision_load_torch_file = _remember_original("clip_vision.load_torch_file", clip_vision.load_torch_file) if clip_vision.load_torch_file is original_clip_vision_load_torch_file: clip_vision.load_torch_file = comfy_utils.load_torch_file _patch_model_management_devices() original_model_patcher_detach = _remember_original("model_patcher.ModelPatcher.detach", model_patcher.ModelPatcher.detach) if model_patcher.ModelPatcher.detach is original_model_patcher_detach: model_patcher.ModelPatcher.detach = _wrap_model_patcher_detach(original_model_patcher_detach) original_load_checkpoint_guess_config = _remember_original( "sd.load_checkpoint_guess_config", comfy_sd.load_checkpoint_guess_config, ) if comfy_sd.load_checkpoint_guess_config is original_load_checkpoint_guess_config: comfy_sd.load_checkpoint_guess_config = _wrap_with_load_context( KIND_CHECKPOINT, path_arg_index=0, bind_output=_bind_checkpoint_outputs, )(original_load_checkpoint_guess_config) original_load_diffusion_model = _remember_original("sd.load_diffusion_model", comfy_sd.load_diffusion_model) if comfy_sd.load_diffusion_model is original_load_diffusion_model: comfy_sd.load_diffusion_model = _wrap_with_load_context( KIND_MODEL, path_arg_index=0, bind_output=lambda model, source_path: model is not None and REGISTRY.bind_object(model, source_path=source_path, kind=KIND_MODEL), )(original_load_diffusion_model) original_load_clip = _remember_original("sd.load_clip", comfy_sd.load_clip) if comfy_sd.load_clip is original_load_clip: comfy_sd.load_clip = _wrap_load_clip(original_load_clip) original_vae_encode = _remember_original("sd.VAE.encode", comfy_sd.VAE.encode) if comfy_sd.VAE.encode is original_vae_encode: comfy_sd.VAE.encode = _wrap_vae_encode(original_vae_encode) original_vae_decode = _remember_original("sd.VAE.decode", comfy_sd.VAE.decode) if comfy_sd.VAE.decode is original_vae_decode: comfy_sd.VAE.decode = _wrap_vae_decode(original_vae_decode) original_clip_vision_load = _remember_original("clip_vision.load", clip_vision.load) if clip_vision.load is original_clip_vision_load: clip_vision.load = _wrap_with_load_context( KIND_CLIP_VISION, path_arg_index=0, bind_output=lambda result, source_path: result is not None and getattr(result, "patcher", None) is not None and REGISTRY.bind_object(result.patcher, source_path=source_path, kind=KIND_CLIP_VISION), )(original_clip_vision_load) original_load_controlnet = _remember_original("controlnet.load_controlnet", controlnet.load_controlnet) if controlnet.load_controlnet is original_load_controlnet: controlnet.load_controlnet = _wrap_with_load_context(KIND_CONTROLNET, path_arg_index=0)(original_load_controlnet) original_load_diffusers = _remember_original("diffusers_load.load_diffusers", diffusers_load.load_diffusers) if diffusers_load.load_diffusers is original_load_diffusers: diffusers_load.load_diffusers = _wrap_with_load_context( KIND_CHECKPOINT, path_arg_index=0, bind_output=_bind_diffusers_outputs, )(original_load_diffusers) REGISTRY.refresh_runtime_state() _PATCHED = True _LOG.info("GPU Resident Loader: monkey patches active on ComfyUI loader and residency paths")