diff --git a/comfy/ops.py b/comfy/ops.py index 34f72ff17..98fec1e1d 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -19,7 +19,7 @@ import torch import logging import comfy.model_management -from comfy.cli_args import args, PerformanceFeature, enables_dynamic_vram +from comfy.cli_args import args, PerformanceFeature import comfy.float import json import comfy.memory_management @@ -296,7 +296,7 @@ class disable_weight_init: class Linear(torch.nn.Linear, CastWeightBiasOp): def __init__(self, in_features, out_features, bias=True, device=None, dtype=None): - if not comfy.model_management.WINDOWS or not enables_dynamic_vram(): + if not comfy.model_management.WINDOWS or not comfy.memory_management.aimdo_enabled: super().__init__(in_features, out_features, bias, device, dtype) return @@ -317,7 +317,7 @@ class disable_weight_init: def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): - if not comfy.model_management.WINDOWS or not enables_dynamic_vram(): + if not comfy.model_management.WINDOWS or not comfy.memory_management.aimdo_enabled: return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) assign_to_params_buffers = local_metadata.get("assign_to_params_buffers", False) diff --git a/comfy/utils.py b/comfy/utils.py index 5fe66ecdb..0769cef44 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -29,7 +29,7 @@ import itertools from torch.nn.functional import interpolate from tqdm.auto import trange from einops import rearrange -from comfy.cli_args import args, enables_dynamic_vram +from comfy.cli_args import args import json import time import mmap @@ -113,7 +113,7 @@ def load_torch_file(ckpt, safe_load=False, device=None, return_metadata=False): metadata = None if ckpt.lower().endswith(".safetensors") or ckpt.lower().endswith(".sft"): try: - if enables_dynamic_vram(): + if comfy.memory_management.aimdo_enabled: sd, metadata = load_safetensors(ckpt) if not return_metadata: metadata = None