diff --git a/nodes/lora_extract_svd.py b/nodes/lora_extract_svd.py index cd3e715..b528ba6 100644 --- a/nodes/lora_extract_svd.py +++ b/nodes/lora_extract_svd.py @@ -13,12 +13,21 @@ import comfy.utils from tqdm import tqdm from comfy_api.latest import io from safetensors.torch import save_file -from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned from .device_utils import ( estimate_model_size, prepare_for_large_operation, cleanup_after_operation, get_device_capabilities ) +try: + from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned + UNIFIED_ENABLED = True +except Exception as e: + UNIFIED_ENABLED = False + +if not UNIFIED_ENABLED: + from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned + + # ============================================================================= # Rank Index Functions - Determine LoRA rank from singular values diff --git a/nodes/lora_merger.py b/nodes/lora_merger.py index 7ef5109..66cd886 100644 --- a/nodes/lora_merger.py +++ b/nodes/lora_merger.py @@ -10,7 +10,6 @@ from tqdm import tqdm from comfy_api.latest import io from safetensors.torch import save_file from typing import List, Dict, Tuple, Optional -from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation from .lora_resize import ( detect_lora_format, @@ -21,6 +20,16 @@ from .lora_resize import ( _format_lora_key ) +try: + from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned + UNIFIED_ENABLED = True +except Exception as e: + UNIFIED_ENABLED = False + +if not UNIFIED_ENABLED: + from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned + + def merge_multi_loras( lora_paths: List[str], lora_weights: List[float], diff --git a/nodes/lora_resize.py b/nodes/lora_resize.py index 07b5876..b75e0d6 100644 --- a/nodes/lora_resize.py +++ b/nodes/lora_resize.py @@ -14,9 +14,18 @@ from tqdm import tqdm from comfy_api.latest import io from safetensors.torch import save_file from typing import Optional, Dict, Tuple, List -from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation +try: + from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned + UNIFIED_ENABLED = True +except Exception as e: + UNIFIED_ENABLED = False + +if not UNIFIED_ENABLED: + from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned + + # Reuse rank functions from extraction module from .lora_extract_svd import ( _index_sv_ratio, diff --git a/nodes/merger.py b/nodes/merger.py index 36de235..3977dd8 100644 --- a/nodes/merger.py +++ b/nodes/merger.py @@ -7,11 +7,18 @@ from tqdm import tqdm from comfy_api.latest import io from safetensors.torch import save_file from .merger_ops import TWO_MODEL_MODES, THREE_MODEL_MODES, MissingTensorBehavior, MissingTensorError -from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned from .device_utils import ( estimate_model_size, prepare_for_large_operation, cleanup_after_operation ) +try: + from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned + UNIFIED_ENABLED = True +except Exception as e: + UNIFIED_ENABLED = False + +if not UNIFIED_ENABLED: + from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned def load_documentation_from_file(filename): """Loads documentation from a markdown file in the ../docs/ directory.""" diff --git a/nodes/prunekeys.py b/nodes/prunekeys.py index 6fbe18f..e0a704e 100644 --- a/nodes/prunekeys.py +++ b/nodes/prunekeys.py @@ -7,9 +7,17 @@ from comfy_api.latest import io from safetensors.torch import save_file from safetensors import safe_open from .utils import convert_pt_to_safetensors -from .merger_utils import MemoryEfficientSafeOpen from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation +try: + from unifiedefficientloader import MemoryEfficientSafeOpen + UNIFIED_ENABLED = True +except Exception as e: + UNIFIED_ENABLED = False + +if not UNIFIED_ENABLED: + from .merger_utils import MemoryEfficientSafeOpen + def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str, use_regex: bool, output_filename: str) -> str: diff --git a/nodes/renamekeys.py b/nodes/renamekeys.py index 625f753..b764591 100644 --- a/nodes/renamekeys.py +++ b/nodes/renamekeys.py @@ -6,9 +6,17 @@ from comfy_api.latest import io from safetensors.torch import save_file from safetensors import safe_open from .utils import convert_pt_to_safetensors -from .merger_utils import MemoryEfficientSafeOpen from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation +try: + from unifiedefficientloader import MemoryEfficientSafeOpen + UNIFIED_ENABLED = True +except Exception as e: + UNIFIED_ENABLED = False + +if not UNIFIED_ENABLED: + from .merger_utils import MemoryEfficientSafeOpen + def _rename_keys(model_name: str, model_type: str, old_keys_str: str, new_keys_str: str, output_filename: str) -> str: