mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-09-04 14:57:08 +08:00
1870 lines
73 KiB
Python
1870 lines
73 KiB
Python
"""
|
|
This file is part of ComfyUI.
|
|
Copyright (C) 2024 Comfy
|
|
This program is free software: you can redistribute it and/or modify
|
|
it under the terms of the GNU General Public License as published by
|
|
the Free Software Foundation, either version 3 of the License, or
|
|
(at your option) any later version.
|
|
This program is distributed in the hope that it will be useful,
|
|
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
GNU General Public License for more details.
|
|
You should have received a copy of the GNU General Public License
|
|
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
"""
|
|
import psutil
|
|
import logging
|
|
import torch
|
|
import sys
|
|
import platform
|
|
import contextlib
|
|
import weakref
|
|
import time
|
|
import gc
|
|
import re
|
|
import threading
|
|
import traceback
|
|
from enum import Enum
|
|
from comfy.cli_args import args, PerformanceFeature
|
|
from comfy.ldm.models.autoencoder import AutoencoderKL
|
|
|
|
try:
|
|
import torch_directml
|
|
_torch_directml_available = True
|
|
except ImportError:
|
|
_torch_directml_available = False
|
|
|
|
def log_vram_state(device=None):
|
|
if not DEBUG_ENABLED:
|
|
return
|
|
if device is None:
|
|
device = get_torch_device()
|
|
free_vram, free_torch = get_free_memory(device, torch_free_too=True)
|
|
active_models = [(m.model.__class__.__name__, m.model_memory_required(device) / 1024**3)
|
|
for m in current_loaded_models if m.device == device]
|
|
logging.debug(
|
|
f"VRAM state: free_vram={free_vram / 1024**3:.2f} GB, free_torch={free_torch / 1024**3:.2f} GB, models={active_models}")
|
|
|
|
class CPUState(Enum):
|
|
GPU = 0
|
|
CPU = 1
|
|
MPS = 2
|
|
|
|
cpu_state = CPUState.GPU # Default to GPU
|
|
|
|
# Global flags
|
|
PROFILING_ENABLED = args.profile
|
|
DEBUG_ENABLED = args.debug
|
|
VERBOSE_ENABLED = False
|
|
|
|
# Configure logging
|
|
logging.basicConfig(level=logging.DEBUG if args.debug or args.profile else logging.INFO)
|
|
|
|
# Cache for device and dtype checks
|
|
_device_cache = {}
|
|
|
|
# VRAM optimizers for extensibility
|
|
_vram_optimizers = []
|
|
|
|
class VRAMState(Enum):
|
|
DISABLED = 0 # No VRAM: models stay on CPU
|
|
NO_VRAM = 1 # Very low VRAM: maximum memory saving
|
|
LOW_VRAM = 2 # Low VRAM: partial model loading
|
|
NORMAL_VRAM = 3 # Default: balanced memory management
|
|
HIGH_VRAM = 4 # High VRAM: keep models in VRAM
|
|
SHARED = 5 # Shared CPU/GPU memory (e.g., MPS)
|
|
|
|
# Global state
|
|
vram_state = VRAMState.NORMAL_VRAM
|
|
set_vram_to = VRAMState.NORMAL_VRAM
|
|
cpu_state = CPUState.GPU
|
|
total_vram = 0
|
|
total_ram = psutil.virtual_memory().total / (1024 * 1024)
|
|
|
|
# Cache for DirectML VRAM
|
|
_directml_vram_cache = {}
|
|
|
|
# Cache for active models memory in DirectML
|
|
_directml_active_memory_cache = {}
|
|
|
|
def cpu_mode():
|
|
"""Check if system is in CPU mode."""
|
|
global cpu_state
|
|
return cpu_state == CPUState.CPU
|
|
|
|
def mps_mode():
|
|
"""Check if system is in MPS (Apple Metal) mode."""
|
|
global cpu_state
|
|
return cpu_state == CPUState.MPS
|
|
|
|
def is_device_cpu(device):
|
|
return is_device_type(device, 'cpu')
|
|
|
|
def is_device_mps(device):
|
|
return is_device_type(device, 'mps')
|
|
|
|
def is_device_cuda(device):
|
|
return is_device_type(device, 'cuda')
|
|
|
|
def is_directml_enabled():
|
|
global directml_enabled
|
|
if directml_enabled:
|
|
return True
|
|
|
|
return False
|
|
|
|
def get_supported_float8_types():
|
|
"""Get supported float8 data types."""
|
|
float8_types = []
|
|
for dtype in [torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2, torch.float8_e5m2fnuz, torch.float8_e8m0fnu]:
|
|
try:
|
|
float8_types.append(dtype)
|
|
except:
|
|
pass
|
|
return float8_types
|
|
|
|
def get_directml_vram(dev):
|
|
"""
|
|
Estimate VRAM for DirectML device, trying CUDA first, then heuristic, then fallback.
|
|
|
|
Args:
|
|
dev: Torch device (DirectML).
|
|
|
|
Returns:
|
|
int: Estimated VRAM in bytes.
|
|
"""
|
|
if dev in _directml_vram_cache:
|
|
return _directml_vram_cache[dev]
|
|
|
|
# Use args.reserve_vram if provided
|
|
if args.reserve_vram is not None:
|
|
vram = int(args.reserve_vram * 1024 * 1024 * 1024)
|
|
_directml_vram_cache[dev] = vram
|
|
return vram
|
|
|
|
# Try CUDA if available
|
|
if torch.cuda.is_available():
|
|
try:
|
|
free_vram, total_vram = torch.cuda.mem_get_info()
|
|
_directml_vram_cache[dev] = total_vram
|
|
if DEBUG_ENABLED:
|
|
logging.debug(f"DirectML VRAM from CUDA: {total_vram / (1024**3):.0f} GB")
|
|
return total_vram
|
|
except Exception as e:
|
|
logging.warning(f"Failed to get CUDA VRAM: {e}")
|
|
|
|
# Try torch_directml heuristic
|
|
if _torch_directml_available:
|
|
try:
|
|
device_index = dev.index if hasattr(dev, 'index') else 0
|
|
device_name = torch_directml.device_name(device_index).lower()
|
|
vram_map = {
|
|
'gtx 1660': 6 * 1024 * 1024 * 1024,
|
|
'gtx 1650': 4 * 1024 * 1024 * 1024,
|
|
'rtx 2060': 6 * 1024 * 1024 * 1024,
|
|
'rtx 3060': 12 * 1024 * 1024 * 1024,
|
|
'rtx 4060': 8 * 1024 * 1024 * 1024,
|
|
'rx 580': 8 * 1024 * 1024 * 1024,
|
|
'rx 570': 8 * 1024 * 1024 * 1024,
|
|
'rx 6700': 12 * 1024 * 1024 * 1024,
|
|
'arc a770': 16 * 1024 * 1024 * 1024,
|
|
}
|
|
vram = 6 * 1024 * 1024 * 1024
|
|
for key, value in vram_map.items():
|
|
if key in device_name:
|
|
vram = value
|
|
break
|
|
_directml_vram_cache[dev] = vram
|
|
if DEBUG_ENABLED:
|
|
logging.debug(f"DirectML VRAM for {device_name}: {vram / (1024**3):.0f} GB")
|
|
return vram
|
|
except Exception as e:
|
|
logging.warning(f"Failed to get DirectML device name: {e}")
|
|
|
|
# Fallback to safe default
|
|
vram = 6 * 1024 * 1024 * 1024
|
|
_directml_vram_cache[dev] = vram
|
|
if DEBUG_ENABLED:
|
|
logging.debug(f"DirectML VRAM fallback: {vram / (1024**3):.0f} GB")
|
|
return vram
|
|
|
|
FLOAT8_TYPES = get_supported_float8_types()
|
|
XFORMERS_IS_AVAILABLE = False
|
|
XFORMERS_ENABLED_VAE = True
|
|
ENABLE_PYTORCH_ATTENTION = True # Enable PyTorch attention for better performance
|
|
FORCE_FP32 = args.force_fp32
|
|
DISABLE_SMART_MEMORY = args.disable_smart_memory
|
|
|
|
# Async offload setup
|
|
STREAMS = {}
|
|
NUM_STREAMS = 1
|
|
stream_counters = {}
|
|
if args.async_offload:
|
|
logging.info(f"Using async weight offloading with {NUM_STREAMS} streams")
|
|
# Protection for older GPUs
|
|
if is_nvidia():
|
|
props = torch.cuda.get_device_properties(get_torch_device())
|
|
if props.major < 8: # Turing (7.5) or Pascal (6.x)
|
|
args.async_offload = False
|
|
NUM_STREAMS = 1
|
|
logging.warning("Async offload disabled for GPUs with SM < 8.0 to prevent memory leaks")
|
|
|
|
# Device initialization
|
|
xpu_available = False
|
|
npu_available = False
|
|
mlu_available = False
|
|
directml_enabled = args.directml is not None
|
|
torch_version_numeric = (0, 0)
|
|
try:
|
|
torch_version = torch.__version__
|
|
temp = torch_version.split(".")
|
|
torch_version_numeric = (int(temp[0]), int(temp[1]))
|
|
xpu_available = (torch_version_numeric[0] < 2 or (torch_version_numeric[0] == 2 and torch_version_numeric[1] <= 4)) and hasattr(
|
|
torch, "xpu") and torch.xpu.is_available()
|
|
except:
|
|
pass
|
|
if directml_enabled:
|
|
import torch_directml
|
|
device_index = args.directml if args.directml >= 0 else 0
|
|
directml_device = torch_directml.device(device_index)
|
|
logging.info(f"Using DirectML with device: {torch_directml.device_name(device_index)}")
|
|
try:
|
|
import intel_extension_for_pytorch as ipex
|
|
xpu_available = xpu_available or torch.xpu.is_available()
|
|
except:
|
|
xpu_available = xpu_available or (hasattr(torch, "xpu") and torch.xpu.is_available())
|
|
try:
|
|
if torch.backends.mps.is_available():
|
|
cpu_state = CPUState.MPS
|
|
import torch.mps
|
|
except:
|
|
pass
|
|
try:
|
|
import torch_npu
|
|
npu_available = torch.npu.is_available()
|
|
except:
|
|
npu_available = False
|
|
try:
|
|
import torch_mlu
|
|
mlu_available = torch.mlu.is_available()
|
|
except:
|
|
mlu_available = False
|
|
if args.cpu:
|
|
cpu_state = CPUState.CPU
|
|
|
|
# Device and memory utilities
|
|
def is_nvidia():
|
|
"""Check if the device is NVIDIA GPU."""
|
|
return cpu_state == CPUState.GPU and torch.version.cuda
|
|
|
|
def is_amd():
|
|
"""Check if the device is AMD GPU."""
|
|
return cpu_state == CPUState.GPU and torch.version.hip
|
|
|
|
def is_intel_xpu():
|
|
"""Check if the device is Intel XPU."""
|
|
return cpu_state == CPUState.GPU and xpu_available
|
|
|
|
def is_ascend_npu():
|
|
"""Check if the device is Ascend NPU."""
|
|
return npu_available
|
|
|
|
def is_mlu():
|
|
"""Check if the device is MLU."""
|
|
return mlu_available
|
|
|
|
def is_device_cuda(device):
|
|
"""Check if the device is CUDA."""
|
|
return hasattr(device, 'type') and device.type == 'cuda'
|
|
|
|
def is_device_type(device, device_type):
|
|
"""Check if the device matches the given type."""
|
|
return hasattr(device, 'type') and device.type == device_type
|
|
|
|
def get_torch_device():
|
|
"""Get the current PyTorch device."""
|
|
if directml_enabled:
|
|
return directml_device
|
|
if cpu_state == CPUState.MPS:
|
|
return torch.device("mps")
|
|
if cpu_state == CPUState.CPU:
|
|
return torch.device("cpu")
|
|
if is_intel_xpu():
|
|
return torch.device("xpu", torch.xpu.current_device())
|
|
if is_ascend_npu():
|
|
return torch.device("npu", torch.npu.current_device())
|
|
if is_mlu():
|
|
return torch.device("mlu", torch.mlu.current_device())
|
|
return torch.device(torch.cuda.current_device())
|
|
|
|
|
|
def get_total_memory(dev=None, torch_total_too=False):
|
|
"""
|
|
Get total memory available on the device.
|
|
|
|
Args:
|
|
dev: Torch device (optional, defaults to current device).
|
|
torch_total_too: If True, return (total, torch_total).
|
|
|
|
Returns:
|
|
int or tuple: Total memory in bytes (or tuple with torch_total).
|
|
"""
|
|
if dev is None:
|
|
dev = get_torch_device()
|
|
if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
|
|
mem_total = psutil.virtual_memory().total
|
|
mem_total_torch = mem_total
|
|
else:
|
|
if directml_enabled:
|
|
mem_total = get_directml_vram(dev)
|
|
mem_total_torch = mem_total
|
|
elif is_intel_xpu():
|
|
stats = torch.xpu.memory_stats(dev)
|
|
mem_reserved = stats['reserved_bytes.all.current']
|
|
mem_total_torch = mem_reserved
|
|
mem_total = torch.xpu.get_device_properties(dev).total_memory
|
|
elif is_ascend_npu():
|
|
stats = torch.npu.memory_stats(dev)
|
|
mem_reserved = stats['reserved_bytes.all.current']
|
|
mem_total_torch = mem_reserved
|
|
_, mem_total_npu = torch.npu.mem_get_info(dev)
|
|
mem_total = mem_total_npu
|
|
elif is_mlu():
|
|
stats = torch.mlu.memory_stats(dev)
|
|
mem_reserved = stats['reserved_bytes.all.current']
|
|
mem_total_torch = mem_reserved
|
|
_, mem_total_mlu = torch.mlu.mem_get_info(dev)
|
|
mem_total = mem_total_mlu
|
|
else:
|
|
stats = torch.cuda.memory_stats(dev)
|
|
mem_reserved = stats['reserved_bytes.all.current']
|
|
_, mem_total_cuda = torch.cuda.mem_get_info(dev)
|
|
mem_total_torch = mem_reserved
|
|
mem_total = mem_total_cuda
|
|
return (mem_total, mem_total_torch) if torch_total_too else mem_total
|
|
|
|
# Initialize VRAM state
|
|
total_vram = get_total_memory(get_torch_device()) / (1024 * 1024)
|
|
logging.info(f"Total VRAM {total_vram:.0f} MB, total RAM {total_ram:.0f} MB")
|
|
logging.info(f"Pytorch version: {torch_version}")
|
|
|
|
def get_extra_reserved_vram():
|
|
"""
|
|
Determine extra VRAM to reserve based on total VRAM and args.
|
|
|
|
Returns:
|
|
int: Reserved VRAM in bytes.
|
|
"""
|
|
total_vram = get_total_memory(get_torch_device()) / (1024 * 1024 * 1024) # VRAM in GB
|
|
if args.reserve_vram is not None:
|
|
return args.reserve_vram * 1024 * 1024 * 1024
|
|
if total_vram < 7.9:
|
|
return 150 * 1024 * 1024 # 150 MB for low VRAM (<7.9 GB)
|
|
return 200 * 1024 * 1024 # 200 MB for high VRAM (?7.9 GB)
|
|
|
|
EXTRA_RESERVED_VRAM = get_extra_reserved_vram()
|
|
logging.info(f"EXTRA_RESERVED_VRAM set to {EXTRA_RESERVED_VRAM / (1024 * 1024):.0f} MB")
|
|
if args.lowvram:
|
|
set_vram_to = VRAMState.LOW_VRAM
|
|
elif args.novram:
|
|
set_vram_to = VRAMState.NO_VRAM
|
|
elif args.highvram or args.gpu_only:
|
|
vram_state = VRAMState.HIGH_VRAM
|
|
if cpu_state != CPUState.GPU:
|
|
vram_state = VRAMState.DISABLED
|
|
elif cpu_state == CPUState.MPS:
|
|
vram_state = VRAMState.SHARED
|
|
if directml_enabled:
|
|
lowvram_available = False
|
|
else:
|
|
lowvram_available = True
|
|
if lowvram_available and set_vram_to in (VRAMState.LOW_VRAM, VRAMState.NO_VRAM):
|
|
vram_state = set_vram_to
|
|
logging.info(f"Set VRAM state to: {vram_state.name}")
|
|
if DISABLE_SMART_MEMORY:
|
|
logging.info("Disabling smart memory management")
|
|
|
|
# XFormers and attention settings
|
|
XFORMERS_VERSION = ""
|
|
if args.disable_xformers:
|
|
XFORMERS_IS_AVAILABLE = False
|
|
else:
|
|
try:
|
|
import xformers
|
|
import xformers.ops
|
|
XFORMERS_IS_AVAILABLE = True
|
|
try:
|
|
XFORMERS_IS_AVAILABLE = xformers._has_cpp_library
|
|
except:
|
|
pass
|
|
XFORMERS_VERSION = xformers.version.__version__
|
|
logging.info(f"xformers version: {XFORMERS_VERSION}")
|
|
if XFORMERS_VERSION.startswith("0.0.18"):
|
|
logging.warning(
|
|
"WARNING: xformers 0.0.18 has a bug causing black images at high resolutions. Please downgrade or upgrade.")
|
|
XFORMERS_ENABLED_VAE = False
|
|
except:
|
|
XFORMERS_IS_AVAILABLE = False
|
|
|
|
def xformers_enabled():
|
|
"""Check if xformers is enabled and available."""
|
|
global directml_enabled, cpu_state
|
|
if cpu_state != CPUState.GPU or is_intel_xpu() or is_ascend_npu() or is_mlu() or directml_enabled:
|
|
return False
|
|
return XFORMERS_IS_AVAILABLE and not args.disable_xformers and not args.use_pytorch_cross_attention
|
|
|
|
def xformers_enabled_vae():
|
|
"""Check if xformers is enabled for VAE."""
|
|
enabled = xformers_enabled()
|
|
if not enabled:
|
|
return False
|
|
return XFORMERS_ENABLED_VAE
|
|
|
|
def sage_attention_enabled():
|
|
"""Check if Sage Attention is enabled."""
|
|
global directml_enabled, cpu_state
|
|
if cpu_state != CPUState.GPU or is_intel_xpu() or is_ascend_npu() or is_mlu() or directml_enabled:
|
|
return False
|
|
return hasattr(args, 'use_sage_attention') and args.use_sage_attention
|
|
|
|
def flash_attention_enabled():
|
|
"""Check if Flash Attention is enabled."""
|
|
global directml_enabled, cpu_state
|
|
if cpu_state != CPUState.GPU or is_intel_xpu() or is_ascend_npu() or is_mlu() or directml_enabled:
|
|
return False
|
|
return hasattr(args, 'use_flash_attention') and args.use_flash_attention
|
|
|
|
def pytorch_attention_enabled():
|
|
"""Check if PyTorch attention is enabled."""
|
|
global ENABLE_PYTORCH_ATTENTION
|
|
return ENABLE_PYTORCH_ATTENTION or not (xformers_enabled() or sage_attention_enabled() or flash_attention_enabled())
|
|
|
|
def pytorch_attention_enabled_vae():
|
|
"""Check if PyTorch attention is enabled for VAE."""
|
|
if is_amd():
|
|
return False # Enabling PyTorch attention on AMD causes crashes at high resolutions
|
|
return pytorch_attention_enabled()
|
|
|
|
def pytorch_attention_flash_attention():
|
|
"""Check if PyTorch Flash Attention is supported."""
|
|
if pytorch_attention_enabled():
|
|
if is_nvidia() or is_intel_xpu() or is_ascend_npu() or is_mlu() or is_amd():
|
|
return True
|
|
return False
|
|
return False
|
|
|
|
def force_upcast_attention_dtype():
|
|
"""Check if attention dtype should be upcast (e.g., FP16 to FP32)."""
|
|
upcast = args.force_upcast_attention
|
|
macos_version = mac_version()
|
|
if macos_version is not None and ((14, 5) <= macos_version < (16,)):
|
|
upcast = True # Workaround for macOS black image bug
|
|
if upcast:
|
|
return {torch.float16: torch.float32}
|
|
return None
|
|
|
|
def cast_to(weight, dtype=None, device=None, non_blocking=False, copy=False, stream=None):
|
|
"""Cast tensor to specified dtype and device, compatible with comfy.ops."""
|
|
if device is None or weight.device == device:
|
|
if not copy:
|
|
if dtype is None or weight.dtype == dtype:
|
|
return weight
|
|
if stream is not None:
|
|
with stream:
|
|
return weight.to(dtype=dtype, copy=copy)
|
|
return weight.to(dtype=dtype, copy=copy)
|
|
if stream is not None:
|
|
with stream:
|
|
r = torch.empty_like(weight, dtype=dtype, device=device)
|
|
r.copy_(weight, non_blocking=non_blocking)
|
|
else:
|
|
r = torch.empty_like(weight, dtype=dtype, device=device)
|
|
r.copy_(weight, non_blocking=non_blocking)
|
|
return r
|
|
|
|
def get_torch_device_name(device=None):
|
|
"""Get the name of the torch device."""
|
|
if device is None:
|
|
device = get_torch_device()
|
|
if isinstance(device, str):
|
|
return device
|
|
if isinstance(device, torch.device):
|
|
if device.type == "cuda":
|
|
try:
|
|
allocator = torch.cuda.get_allocator_backend()
|
|
except:
|
|
allocator = ""
|
|
return f"{device.type}:{device.index if device.index is not None else 0} {allocator}"
|
|
return device.type
|
|
return str(device)
|
|
|
|
class OOM_EXCEPTION(Exception):
|
|
"""Exception raised for out-of-memory errors."""
|
|
pass
|
|
|
|
if args.use_pytorch_cross_attention:
|
|
ENABLE_PYTORCH_ATTENTION = True
|
|
XFORMERS_IS_AVAILABLE = False
|
|
MIN_WEIGHT_MEMORY_RATIO = 0.4 if is_nvidia() else 0.0
|
|
if is_nvidia() and torch_version_numeric[0] >= 2:
|
|
if not (ENABLE_PYTORCH_ATTENTION or args.use_split_cross_attention or args.use_quad_cross_attention):
|
|
ENABLE_PYTORCH_ATTENTION = True
|
|
elif is_intel_xpu() or is_ascend_npu() or is_mlu():
|
|
if not (args.use_split_cross_attention or args.use_quad_cross_attention):
|
|
ENABLE_PYTORCH_ATTENTION = True
|
|
elif is_amd() and torch_version_numeric[0] >= 2 and torch_version_numeric[1] >= 7:
|
|
arch = torch.cuda.get_device_properties(get_torch_device()).gcnArchName
|
|
logging.info(f"AMD arch: {arch}")
|
|
if any(a in arch for a in ["gfx1100", "gfx1101"]) and not (args.use_split_cross_attention or args.use_quad_cross_attention):
|
|
ENABLE_PYTORCH_ATTENTION = True
|
|
if ENABLE_PYTORCH_ATTENTION:
|
|
torch.backends.cuda.enable_math_sdp(True)
|
|
torch.backends.cuda.enable_flash_sdp(True)
|
|
torch.backends.cuda.enable_mem_efficient_sdp(True)
|
|
if torch_version_numeric[0] == 2 and torch_version_numeric[1] >= 5:
|
|
torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(True)
|
|
else:
|
|
logging.warning("Could not set allow_fp16_bf16_reduction_math_sdp")
|
|
|
|
def get_free_memory(dev=None, torch_free_too=False):
|
|
"""
|
|
Get free memory available on the device.
|
|
|
|
Args:
|
|
dev: Torch device (optional, defaults to current device).
|
|
torch_free_too: If True, return (free_total, free_torch).
|
|
|
|
Returns:
|
|
int or tuple: Free memory in bytes (or tuple with free_torch).
|
|
"""
|
|
if dev is None:
|
|
dev = get_torch_device()
|
|
if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
|
|
mem_free_total = psutil.virtual_memory().available
|
|
mem_free_torch = mem_free_total
|
|
else:
|
|
if directml_enabled:
|
|
total_vram = get_directml_vram(dev)
|
|
cache_key = (dev, 'active_models')
|
|
if cache_key not in _directml_active_memory_cache:
|
|
active_models = sum(m.model_loaded_memory() for m in current_loaded_models if m.device == dev)
|
|
_directml_active_memory_cache[cache_key] = active_models
|
|
active_models = _directml_active_memory_cache[cache_key]
|
|
mem_free_total = max(1024 * 1024 * 1024, total_vram - active_models * 1.2)
|
|
mem_free_torch = mem_free_total
|
|
if DEBUG_ENABLED:
|
|
logging.debug(f"DirectML: total_vram={total_vram / (1024**3):.0f} GB, active_models={active_models / (1024**3):.2f} GB, free={mem_free_total / (1024**3):.2f} GB")
|
|
elif is_intel_xpu():
|
|
stats = torch.xpu.memory_stats(dev)
|
|
mem_active = stats['active_bytes.all.current']
|
|
mem_reserved = stats['reserved_bytes.all.current']
|
|
mem_free_torch = mem_reserved - mem_active
|
|
mem_free_xpu = torch.xpu.get_device_properties(dev).total_memory - mem_reserved
|
|
mem_free_total = mem_free_xpu + mem_free_torch
|
|
elif is_ascend_npu():
|
|
stats = torch.npu.memory_stats(dev)
|
|
mem_active = stats['active_bytes.all.current']
|
|
mem_reserved = stats['reserved_bytes.all.current']
|
|
mem_free_npu, _ = torch.npu.mem_get_info(dev)
|
|
mem_free_torch = mem_reserved - mem_active
|
|
mem_free_total = mem_free_npu + mem_free_torch
|
|
elif is_mlu():
|
|
stats = torch.mlu.memory_stats(dev)
|
|
mem_active = stats['active_bytes.all.current']
|
|
mem_reserved = stats['reserved_bytes.all.current']
|
|
mem_free_mlu, _ = torch.mlu.mem_get_info(dev)
|
|
mem_free_torch = mem_reserved - mem_active
|
|
mem_free_total = mem_free_mlu + mem_free_torch
|
|
else:
|
|
stats = torch.cuda.memory_stats(dev)
|
|
mem_active = stats['active_bytes.all.current']
|
|
mem_reserved = stats['reserved_bytes.all.current']
|
|
mem_free_cuda, _ = torch.cuda.mem_get_info(dev)
|
|
mem_free_torch = mem_reserved - mem_active
|
|
mem_free_total = mem_free_cuda + mem_free_torch
|
|
return (mem_free_total, mem_free_torch) if torch_free_too else mem_free_total
|
|
|
|
def get_adaptive_min_free(mem_total, memory_required=None):
|
|
"""
|
|
Calculate adaptive min_free VRAM based on GPU memory and model requirements.
|
|
|
|
Args:
|
|
mem_total (float): Total GPU VRAM in GB.
|
|
memory_required (float, optional): Estimated memory required by the model in GB.
|
|
|
|
Returns:
|
|
float: Minimum free VRAM required in GB.
|
|
"""
|
|
# Base min_free as a fraction of total VRAM
|
|
base_min_free = mem_total * 0.25 # 25% of total VRAM as baseline
|
|
|
|
if memory_required is not None:
|
|
min_free = max(base_min_free, memory_required) # Use memory_required directly, no extra multiplier
|
|
else:
|
|
min_free = base_min_free
|
|
|
|
# Cap min_free to avoid excessive requirements
|
|
min_free = min(min_free, mem_total * 0.5) # Never exceed 50% of total VRAM
|
|
|
|
# Minimum threshold for very small GPUs
|
|
min_free = max(min_free, 1.0 if mem_total < 6.0 else 1.5)
|
|
|
|
if PROFILING_ENABLED:
|
|
memory_required_str = f"{memory_required:.2f}" if memory_required is not None else "None"
|
|
logging.debug(f"get_adaptive_min_free: mem_total={mem_total:.2f} GB, memory_required={memory_required_str} GB, min_free={min_free:.2f} GB")
|
|
|
|
return min_free
|
|
|
|
def memory_monitor(device, interval=5.0):
|
|
"""Monitor memory usage in a background thread."""
|
|
if not DEBUG_ENABLED:
|
|
return
|
|
|
|
def monitor():
|
|
while True:
|
|
log_vram_state(device)
|
|
time.sleep(interval)
|
|
threading.Thread(target=monitor, daemon=True).start()
|
|
|
|
memory_monitor(get_torch_device())
|
|
|
|
def soft_empty_cache(clear=False, device=None, caller="unknown"):
|
|
"""
|
|
Clear PyTorch memory cache efficiently with VRAM check.
|
|
|
|
Args:
|
|
clear (bool): Force cache clearing regardless of memory state.
|
|
device (torch.device): Device to clear cache for. Defaults to current device.
|
|
caller (str): Source of the call for debugging.
|
|
"""
|
|
if device is None:
|
|
device = get_torch_device()
|
|
|
|
if PROFILING_ENABLED:
|
|
start_time = time.time()
|
|
logging.debug(f"soft_empty_cache called with clear={clear}, device={device}, caller={caller}")
|
|
|
|
# Fixed threshold in bytes (100 MB)
|
|
MEMORY_THRESHOLD = 100 * 1024 * 1024 # 100 MB
|
|
cache_key = (device, 'free_memory')
|
|
|
|
mem_free_total, mem_free_torch = get_free_memory(device, torch_free_too=True)
|
|
total_vram = get_total_memory(device)
|
|
if not clear and (mem_free_torch <= MEMORY_THRESHOLD or mem_free_total > 0.4 * total_vram):
|
|
if PROFILING_ENABLED:
|
|
logging.debug(f"soft_empty_cache: Skipped (free_vram={mem_free_total/1024**3:.2f} GB, free_torch={mem_free_torch/1024**3:.2f} GB)")
|
|
return
|
|
|
|
try:
|
|
if is_device_cuda(device):
|
|
torch.cuda.empty_cache()
|
|
if clear and mem_free_torch < min(0.1 * total_vram, 1.0 * 1024**3):
|
|
gc.collect()
|
|
if torch.distributed.is_initialized():
|
|
torch.cuda.ipc_collect()
|
|
elif cpu_state == CPUState.MPS:
|
|
torch.mps.empty_cache()
|
|
elif is_intel_xpu():
|
|
torch.xpu.empty_cache()
|
|
elif is_ascend_npu():
|
|
torch.npu.empty_cache()
|
|
elif is_mlu():
|
|
torch.mlu.empty_cache()
|
|
|
|
if PROFILING_ENABLED:
|
|
free_vram_after, free_torch_after = get_free_memory(device, torch_free_too=True)
|
|
logging.debug(f"After clear: free_vram={free_vram_after/1024**3:.2f} GB, free_torch={free_torch_after/1024**3:.2f} GB")
|
|
logging.debug(f"soft_empty_cache took {time.time() - start_time:.3f} s, gained={(free_vram_after - mem_free_total)/1024**3:.2f} GB")
|
|
except Exception as e:
|
|
if PROFILING_ENABLED:
|
|
logging.warning(f"Failed to clear cache for {device}: {str(e)}")
|
|
|
|
def unload_all_models():
|
|
"""
|
|
Unload all models from memory and clear cache if necessary.
|
|
"""
|
|
if PROFILING_ENABLED:
|
|
start_time = time.time()
|
|
logging.debug("unload_all_models called")
|
|
|
|
for model in list(current_loaded_models):
|
|
model.model_unload()
|
|
|
|
current_loaded_models.clear()
|
|
|
|
device = get_torch_device()
|
|
free_vram = get_free_memory(device)[0]
|
|
total_vram = get_total_memory(device)
|
|
clear_aggressive = free_vram < 0.4 * total_vram
|
|
if PROFILING_ENABLED:
|
|
logging.debug(f"unload_all_models: free_vram={free_vram/1024**3:.2f} GB, aggressive_clear={clear_aggressive}")
|
|
|
|
soft_empty_cache(clear=clear_aggressive, caller="unload_all_models")
|
|
|
|
if PROFILING_ENABLED:
|
|
new_free_vram = get_free_memory(device)[0]
|
|
logging.debug(f"unload_all_models done: free={new_free_vram/1024**3:.2f} GB, "
|
|
f"gained={(new_free_vram - free_vram)/1024**3:.2f} GB, took={time.time() - start_time:.3f} s")
|
|
|
|
def get_offload_stream(device):
|
|
"""Get a stream for asynchronous weight offloading."""
|
|
stream_counter = stream_counters.get(device, 0)
|
|
if NUM_STREAMS <= 1 or not is_device_cuda(device):
|
|
return None
|
|
if device in STREAMS:
|
|
ss = STREAMS[device]
|
|
s = ss[stream_counter]
|
|
stream_counter = (stream_counter + 1) % len(ss)
|
|
if is_device_cuda(device):
|
|
ss[stream_counter].wait_stream(torch.cuda.current_stream())
|
|
stream_counters[device] = stream_counter
|
|
return s
|
|
elif is_device_cuda(device):
|
|
ss = [torch.cuda.Stream(device=device, priority=0) for _ in range(NUM_STREAMS)]
|
|
STREAMS[device] = ss
|
|
s = ss[stream_counter]
|
|
stream_counter = (stream_counter + 1) % len(ss)
|
|
stream_counters[device] = stream_counter
|
|
return s
|
|
return None
|
|
|
|
def sync_stream(device, stream):
|
|
"""Synchronize the given stream with the current CUDA stream."""
|
|
if stream is None or not is_device_cuda(device):
|
|
return
|
|
torch.cuda.current_stream().wait_stream(stream)
|
|
|
|
def cast_to_device(tensor, device, dtype, copy=False):
|
|
"""Cast tensor to specified device and dtype with non-blocking support."""
|
|
non_blocking = device_supports_non_blocking(device)
|
|
return cast_to(tensor, dtype=dtype, device=device, non_blocking=non_blocking, copy=copy)
|
|
|
|
def register_vram_optimizer(optimizer):
|
|
"""Register a VRAM optimizer."""
|
|
_vram_optimizers.append(optimizer)
|
|
|
|
# Model management
|
|
current_loaded_models = []
|
|
|
|
class LoadedModel:
|
|
def __init__(self, model):
|
|
self._set_model(model)
|
|
self.device = model.load_device
|
|
self.real_model = None
|
|
self.currently_used = True
|
|
self.model_offloaded = False
|
|
self.model_finalizer = None
|
|
self._patcher_finalizer = None
|
|
|
|
def _set_model(self, model):
|
|
self._model = weakref.ref(model)
|
|
if hasattr(model, 'parent') and model.parent is not None:
|
|
self._parent_model = weakref.ref(model.parent)
|
|
self._patcher_finalizer = weakref.finalize(model, self._switch_parent)
|
|
|
|
def _switch_parent(self):
|
|
if hasattr(self, '_parent_model'):
|
|
model = self._parent_model()
|
|
if model is not None:
|
|
self._set_model(model)
|
|
|
|
@property
|
|
def model(self):
|
|
return self._model()
|
|
|
|
def model_memory(self):
|
|
return self.model.model_size() if hasattr(self.model, 'model_size') else module_size(self.model)
|
|
|
|
def model_loaded_memory(self):
|
|
return self.model.loaded_size() if hasattr(self.model, 'loaded_size') else module_size(self.model)
|
|
|
|
def model_offloaded_memory(self):
|
|
return self.model_memory() - self.model_loaded_memory()
|
|
|
|
def model_memory_required(self, device):
|
|
"""
|
|
Estimate memory required for the model on the specified device.
|
|
|
|
Args:
|
|
device (torch.device): Target device for memory estimation.
|
|
|
|
Returns:
|
|
int: Memory required in bytes.
|
|
"""
|
|
# Fast path: use size if available
|
|
if hasattr(self.model, 'size') and self.model.size > 0:
|
|
return self.model.size
|
|
|
|
# Check if model is already on the target device
|
|
if hasattr(self.model, 'current_loaded_device') and device == self.model.current_loaded_device():
|
|
return self.model_offloaded_memory()
|
|
|
|
# Handle AutoencoderKL
|
|
if self.model is not None and isinstance(self.model.model, AutoencoderKL):
|
|
shape = getattr(self.model, 'last_shape', (1, 4, 64, 64))
|
|
dtype = getattr(self.model, 'model_dtype', torch.float32)()
|
|
return estimate_vae_decode_memory(self.model.model, shape, dtype)
|
|
|
|
# Sum memory for additional models
|
|
loaded_memory = 0
|
|
if hasattr(self.model, 'additional_models'):
|
|
model_device = device
|
|
if hasattr(self.model.model, 'device'):
|
|
model_device = self.model.model.device
|
|
if DEBUG_ENABLED:
|
|
logging.debug(f"[DEBUG_CLONES] Model {self.model.__class__.__name__} using device {model_device}")
|
|
for m in self.model.additional_models:
|
|
try:
|
|
loaded_memory += m.model_memory_required(model_device)
|
|
except Exception as e:
|
|
if DEBUG_ENABLED:
|
|
logging.warning(f"[DEBUG_CLONES] Error calculating memory for additional model: {e}")
|
|
|
|
return self.model_memory() + loaded_memory
|
|
|
|
def model_load(self, lowvram_model_memory=0, force_patch_weights=False):
|
|
with profile_section("Model load"):
|
|
self.model.model_patches_to(self.device)
|
|
self.model.model_patches_to(self.model.model_dtype())
|
|
use_more_vram = lowvram_model_memory if lowvram_model_memory > 0 else float('inf')
|
|
self.model_use_more_vram(use_more_vram, force_patch_weights=force_patch_weights)
|
|
real_model = self.model.model
|
|
if is_intel_xpu() and not args.disable_ipex_optimize and 'ipex' in globals():
|
|
with torch.no_grad():
|
|
real_model = ipex.optimize(real_model.eval(), inplace=True, graph_mode=True, concat_linear=True)
|
|
self.real_model = weakref.ref(real_model)
|
|
self.model_finalizer = weakref.finalize(real_model, cleanup_models)
|
|
return real_model
|
|
|
|
def should_reload_model(self, force_patch_weights=False):
|
|
return force_patch_weights and self.model.lowvram_patch_counter() > 0
|
|
|
|
def model_unload(self, memory_to_free=None, unpatch_weights=True):
|
|
"""
|
|
Unload the model, freeing memory on both CPU and GPU.
|
|
Clears CUDA cache if needed, logs critical information.
|
|
|
|
Args:
|
|
memory_to_free: Amount of memory to free (bytes), if partial unloading is needed.
|
|
unpatch_weights: Whether to unpatch model weights during unloading.
|
|
|
|
Returns:
|
|
float: Estimated memory freed (in bytes).
|
|
"""
|
|
with profile_section("Model unload"):
|
|
if self.is_dead() or self.real_model is None:
|
|
if DEBUG_ENABLED:
|
|
logging.debug("[DEBUG_CLONES] Model is dead or real_model is None, skipping unload")
|
|
return 0
|
|
|
|
mem_freed = getattr(self.model, 'model_loaded_weight_memory', 0) if self.model is not None else 0
|
|
is_cuda = is_device_cuda(self.device)
|
|
|
|
try:
|
|
model_name = self.model.__class__.__name__ if self.model is not None else "None"
|
|
model_type = self.model.model.__class__.__name__ if self.model is not None and hasattr(self.model, 'model') else "Unknown"
|
|
if DEBUG_ENABLED:
|
|
logging.debug(f"[DEBUG_CLONES] Starting unload for {model_name}(type={model_type})")
|
|
|
|
# Partial unload if requested and supported
|
|
if memory_to_free is not None and memory_to_free < mem_freed and hasattr(self.model, 'partially_unload'):
|
|
freed = self.model.partially_unload(self.model.offload_device, memory_to_free)
|
|
if freed >= memory_to_free:
|
|
if PROFILING_ENABLED:
|
|
logging.debug(f"[DEBUG_CLONES] Partial unload freed {freed / 1024**3:.2f} GB")
|
|
return freed
|
|
|
|
# Full unload
|
|
if self.model is not None and hasattr(self.model, 'detach'):
|
|
self.model.detach(unpatch_all=unpatch_weights)
|
|
if self.model_finalizer is not None:
|
|
self.model_finalizer.detach()
|
|
self.model_finalizer = None
|
|
self.real_model = None
|
|
self._model = lambda: None
|
|
|
|
# Garbage collection for non-CUDA devices
|
|
if not is_cuda:
|
|
gc.collect()
|
|
|
|
# Clear CUDA cache if on CUDA device
|
|
if is_cuda:
|
|
device = self.device
|
|
free_vram = get_free_memory(device)[0]
|
|
total_vram = get_total_memory(device)
|
|
clear_aggressive = free_vram < 0.4 * total_vram
|
|
soft_empty_cache(clear=clear_aggressive, caller="model_unload")
|
|
|
|
if PROFILING_ENABLED:
|
|
logging.debug(f"[DEBUG_CLONES] Unload complete for {model_name}")
|
|
return mem_freed
|
|
|
|
except Exception as e:
|
|
if DEBUG_ENABLED:
|
|
logging.warning(f"[DEBUG_CLONES] Error during model_unload for {model_name}(type={model_type}): {e}")
|
|
return mem_freed
|
|
|
|
def model_use_more_vram(self, use_more_vram, force_patch_weights=False):
|
|
if not use_more_vram:
|
|
if PROFILING_ENABLED:
|
|
logging.debug(
|
|
"model_use_more_vram: use_more_vram=False, returning 0")
|
|
return 0
|
|
mem_required = self.model_memory_required(self.device)
|
|
extra_memory = min(mem_required * 0.3, 50 * 1024 * 1024 * 1024) # Reduced to 50 MB chunks
|
|
return self.model.partially_load(self.device, extra_memory, force_patch_weights=force_patch_weights)
|
|
|
|
def __eq__(self, other):
|
|
return self.model is other.model
|
|
|
|
def __del__(self):
|
|
if hasattr(self, '_patcher_finalizer') and self._patcher_finalizer is not None:
|
|
self._patcher_finalizer.detach()
|
|
if hasattr(self, '_model_finalizer') and self._model_finalizer is not None:
|
|
self._model_finalizer.detach()
|
|
|
|
def is_dead(self):
|
|
"""
|
|
Check if the model is dead (real_model exists but model is garbage collected).
|
|
Returns True if the model is dead, False otherwise.
|
|
"""
|
|
if self.real_model is None:
|
|
return False # Model was never loaded or already unloaded
|
|
return self.real_model() is not None and self.model is None
|
|
|
|
def module_size(model, shape=None, dtype=None):
|
|
"""
|
|
Estimate memory size of a module by summing parameter and buffer sizes,
|
|
or using VAE-specific estimation if shape and dtype are provided.
|
|
"""
|
|
from diffusers import AutoencoderKL
|
|
|
|
module_mem = 0
|
|
if shape is not None and dtype is not None and isinstance(model, AutoencoderKL):
|
|
try:
|
|
batch, channels, height, width = shape
|
|
# Adjusted memory estimate for VAE: reduced multiplier from 64*1.1 to 32*1.05 to avoid overestimation
|
|
base_memory = height * width * channels * 32 * 1.05
|
|
size_of_dtype = dtype_size(dtype)
|
|
module_mem = base_memory * size_of_dtype
|
|
# Add parameter memory for VAE to account for model weights
|
|
param_mem = sum(p.numel() * p.element_size() for p in model.parameters())
|
|
module_mem += param_mem
|
|
if DEBUG_ENABLED:
|
|
logging.debug(
|
|
f"Estimated VAE memory: shape={shape}, dtype={dtype}, "
|
|
f"params={param_mem / (1024**3):.2f} GB, total={module_mem / (1024**3):.2f} GB"
|
|
)
|
|
except Exception as e:
|
|
logging.warning(f"Failed to estimate VAE memory for {model.__class__.__name__}: {str(e)}")
|
|
|
|
if module_mem == 0:
|
|
try:
|
|
# Sum memory of state dict (parameters and buffers)
|
|
module_mem = sum(p.numel() * p.element_size() for p in model.state_dict().values())
|
|
except AttributeError:
|
|
# Fallback: sum parameters and buffers separately
|
|
if hasattr(model, 'parameters'):
|
|
module_mem += sum(p.numel() * p.element_size() for p in model.parameters())
|
|
if hasattr(model, 'buffers'):
|
|
module_mem += sum(b.numel() * b.element_size() for b in model.buffers())
|
|
if module_mem == 0:
|
|
model_name = model.__class__.__name__.lower()
|
|
if 'vae' in model_name or isinstance(model, AutoencoderKL):
|
|
# Reduced fallback from 3.5 GB to 2.5 GB for VAE
|
|
module_mem = 2.5 * 1024**3
|
|
logging.warning(
|
|
f"Could not estimate module size for {model.__class__.__name__}, "
|
|
f"assuming 2.5 GB for VAE"
|
|
)
|
|
else:
|
|
# Minimal memory assumption for unknown models
|
|
module_mem = 1024 * 1024
|
|
logging.warning(
|
|
f"Could not estimate module size for {model.__class__.__name__}, "
|
|
f"assuming minimal memory (1 MB)"
|
|
)
|
|
|
|
if VERBOSE_ENABLED:
|
|
logging.debug(f"Module size for {model.__class__.__name__}: {module_mem / (1024**3):.2f} GB")
|
|
return module_mem
|
|
|
|
def dtype_size(dtype):
|
|
"""Get the size of a data type in bytes."""
|
|
dtype_size = 4
|
|
if dtype in (torch.float16, torch.bfloat16):
|
|
dtype_size = 2
|
|
elif dtype == torch.float32:
|
|
dtype_size = 4
|
|
elif dtype in FLOAT8_TYPES:
|
|
dtype_size = 1
|
|
else:
|
|
try:
|
|
dtype_size = dtype.itemsize
|
|
except:
|
|
pass
|
|
return dtype_size
|
|
|
|
def get_adaptive_buffer(device):
|
|
"""
|
|
Calculate adaptive memory buffer based on VRAM size and available memory.
|
|
|
|
Args:
|
|
device: Device to calculate buffer for.
|
|
|
|
Returns:
|
|
Buffer size in bytes.
|
|
"""
|
|
mem_total = get_total_memory(device) if is_device_cuda(device) else 6 * 1024**3
|
|
mem_free_total, _ = get_free_memory(device, torch_free_too=True)
|
|
# Use 2% for low VRAM (<50% free) or small GPUs, 5% otherwise
|
|
fraction = 0.02 if mem_free_total < mem_total * 0.5 or mem_total < 8 * 1024**3 else 0.05
|
|
buffer = min(max(fraction * mem_total, 0.05 * 1024**3), 0.2 * 1024**3) # 0.05-0.2 GB
|
|
if PROFILING_ENABLED:
|
|
logging.debug(f"get_adaptive_buffer: mem_total={mem_total / 1024**3:.2f} GB, "
|
|
f"mem_free_total={mem_free_total / 1024**3:.2f} GB, buffer={buffer / 1024**3:.2f} GB")
|
|
return buffer
|
|
|
|
def estimate_vae_decode_memory(model, shape, dtype):
|
|
"""
|
|
Estimate memory required for VAE decoding.
|
|
Uses module_size with shape and dtype for accurate estimation.
|
|
"""
|
|
total_memory = module_size(model, shape=shape, dtype=dtype)
|
|
if PROFILING_ENABLED:
|
|
logging.debug(
|
|
f"Estimated VAE decode memory: shape={shape}, dtype={dtype}, "
|
|
f"total={total_memory / (1024**3):.2f} GB"
|
|
)
|
|
return total_memory
|
|
|
|
|
|
def use_more_memory(extra_memory, loaded_models, device):
|
|
"""Use additional VRAM for loaded models."""
|
|
for m in loaded_models:
|
|
if m.device == device:
|
|
extra_memory -= m.model_use_more_vram(extra_memory)
|
|
if extra_memory <= 0:
|
|
break
|
|
|
|
def offloaded_memory(loaded_models, device):
|
|
"""Calculate offloaded memory for loaded models."""
|
|
offloaded_mem = 0
|
|
for m in loaded_models:
|
|
if m.device == device:
|
|
offloaded_mem += m.model_offloaded_memory()
|
|
return offloaded_mem
|
|
|
|
def extra_reserved_memory():
|
|
"""Get extra reserved VRAM."""
|
|
return EXTRA_RESERVED_VRAM
|
|
|
|
def minimum_inference_memory():
|
|
"""Get minimum memory required for inference."""
|
|
return (1024 * 1024 * 1024) * 0.6 + extra_reserved_memory() # Reduced to 600 MB
|
|
|
|
|
|
def cleanup_models_gc():
|
|
"""Clean up dead models and collect garbage if significant memory is freed."""
|
|
dead_memory = 0
|
|
for cur in current_loaded_models:
|
|
if cur.is_dead():
|
|
dead_memory += cur.model_memory()
|
|
|
|
if dead_memory > 50 * 1024 * 1024: # 50 MB threshold
|
|
if PROFILING_ENABLED:
|
|
device = get_torch_device()
|
|
free_vram = get_free_memory(device)[0]
|
|
total_vram = get_total_memory(device)
|
|
logging.debug(f"cleanup_models_gc: dead_memory={dead_memory/1024**2:.2f} MB, "
|
|
f"free_vram={free_vram/1024**3:.2f} GB")
|
|
|
|
soft_empty_cache(clear=False, caller="cleanup_models_gc")
|
|
|
|
i = len(current_loaded_models) - 1
|
|
while i >= 0:
|
|
if current_loaded_models[i].is_dead():
|
|
logging.warning(f"Removing dead model {current_loaded_models[i].real_model().__class__.__name__}")
|
|
current_loaded_models.pop(i)
|
|
i -= 1
|
|
|
|
def free_memory(memory_required, device, keep_loaded=None, loaded_models=None, caller="unknown"):
|
|
"""
|
|
Free memory on the device by unloading models efficiently, prioritizing unused models.
|
|
|
|
Args:
|
|
memory_required (int): Memory needed in bytes.
|
|
device (torch.device): Device to free memory on.
|
|
keep_loaded (list, optional): Models to keep loaded. Defaults to [].
|
|
loaded_models (list, optional): List of models to consider for unloading. Defaults to current_loaded_models.
|
|
caller (str): Source of the call for debugging.
|
|
|
|
Returns:
|
|
list: Unloaded models.
|
|
"""
|
|
with profile_section("free_memory"):
|
|
# Initialize defaults
|
|
if keep_loaded is None:
|
|
keep_loaded = []
|
|
if loaded_models is None:
|
|
loaded_models = current_loaded_models
|
|
|
|
# Cache memory state to avoid redundant calls
|
|
cache_key = (device, 'free_memory')
|
|
mem_free = _device_cache.get(cache_key, None)
|
|
if mem_free is None:
|
|
mem_free = get_free_memory(device, torch_free_too=True)
|
|
_device_cache[cache_key] = mem_free
|
|
mem_free_total = mem_free[0] if isinstance(mem_free, tuple) else mem_free
|
|
total_vram = get_total_memory(device)
|
|
|
|
# Log initial state if profiling is enabled
|
|
if PROFILING_ENABLED:
|
|
logging.debug(
|
|
f"free_memory: requested={memory_required / 1024**3:.2f} GB, "
|
|
f"free={mem_free_total / 1024**3:.2f} GB, models={len(loaded_models)}, "
|
|
f"device={device}, caller={caller}"
|
|
)
|
|
|
|
# Skip if enough VRAM (>20% of total or required memory available)
|
|
if mem_free_total > max(memory_required, 0.4 * total_vram):
|
|
if PROFILING_ENABLED:
|
|
logging.debug(f"free_memory: Skipped (free_vram={mem_free_total / 1024**3:.2f} GB)")
|
|
return []
|
|
|
|
# Apply VRAM optimizers
|
|
for optimizer in _vram_optimizers:
|
|
memory_required = optimizer(memory_required, device, keep_loaded)
|
|
|
|
# Ensure minimum inference memory
|
|
memory_required = max(memory_required, minimum_inference_memory())
|
|
|
|
# Clean up dead models
|
|
cleanup_models_gc()
|
|
|
|
unloaded_models = []
|
|
can_unload = []
|
|
|
|
# Collect models that can be unloaded
|
|
for i in range(len(loaded_models) - 1, -1, -1):
|
|
model = loaded_models[i]
|
|
if model.device == device and model not in keep_loaded and not model.is_dead() and not model.currently_used:
|
|
mem_required = model.model_memory_required(device)
|
|
can_unload.append((mem_required, model, i))
|
|
can_unload.sort(reverse=True) # Prioritize models using more memory
|
|
|
|
# Calculate memory to free
|
|
memory_to_free = memory_required - mem_free_total + extra_reserved_memory()
|
|
|
|
# Unload models to free required memory
|
|
for mem, model, index in can_unload:
|
|
try:
|
|
model_id = getattr(model.model, 'model_id', id(model.model) if hasattr(model.model, 'model') else model.model.__class__.__name__)
|
|
model_type = model.model.__class__.__name__ if hasattr(model.model, 'model') else 'Unknown'
|
|
mem_freed = model.model_unload(memory_to_free=memory_to_free)
|
|
loaded_models.pop(index)
|
|
unloaded_models.append(model)
|
|
if model.model is not None and hasattr(model.model, 'detach'):
|
|
model.model.detach(unpatch_all=True)
|
|
mem_free_total += mem_freed
|
|
_device_cache[cache_key] = (mem_free_total, mem_free[1] if isinstance(mem_free, tuple) else 0)
|
|
if PROFILING_ENABLED:
|
|
logging.debug(
|
|
f"Unloaded model: id={model_id}, type={model_type}, "
|
|
f"freed={mem_freed / 1024**3:.2f} GB, free_vram={mem_free_total / 1024**3:.2f} GB"
|
|
)
|
|
if mem_free_total >= memory_required:
|
|
break
|
|
except Exception as e:
|
|
if DEBUG_ENABLED:
|
|
logging.warning(f"Failed to unload model at index {index}: {e}")
|
|
|
|
# Utilize excess memory if available
|
|
use_more_memory(mem_free_total - memory_required, loaded_models, device)
|
|
|
|
if PROFILING_ENABLED:
|
|
logging.debug(
|
|
f"free_memory done: free={mem_free_total / 1024**3:.2f} GB, unloaded={len(unloaded_models)} models"
|
|
)
|
|
|
|
return unloaded_models
|
|
|
|
def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimum_memory_required=None, force_full_load=False):
|
|
"""
|
|
Load multiple models to GPU, managing VRAM efficiently.
|
|
|
|
Args:
|
|
models: List of models to load.
|
|
memory_required: Estimated memory needed (bytes).
|
|
force_patch_weights: Force re-patching model weights.
|
|
minimum_memory_required: Minimum memory needed for inference.
|
|
force_full_load: Force full model loading regardless of VRAM state.
|
|
"""
|
|
cleanup_models_gc()
|
|
with profile_section("load_models_gpu"):
|
|
# Memory cache for efficient memory queries
|
|
memory_cache = {}
|
|
def get_cached_memory(device, torch_free_too=False):
|
|
cache_key = (device, torch_free_too)
|
|
if cache_key not in memory_cache:
|
|
try:
|
|
memory_cache[cache_key] = get_free_memory(device, torch_free_too)
|
|
except Exception as e:
|
|
logging.error(f"Failed to get memory for {device}: {e}")
|
|
memory_cache[cache_key] = (0, 0) if torch_free_too else 0
|
|
return memory_cache[cache_key]
|
|
|
|
if minimum_memory_required is None:
|
|
minimum_memory_required = minimum_inference_memory()
|
|
device = get_torch_device()
|
|
if vram_state in (VRAMState.DISABLED, VRAMState.SHARED):
|
|
return
|
|
|
|
model_lookup = {m.model: m for m in current_loaded_models if m.model is not None}
|
|
|
|
# Reset currently_used flag for all loaded models
|
|
for loaded_model in current_loaded_models:
|
|
loaded_model.currently_used = False
|
|
|
|
loaded = []
|
|
# Prepare models to load
|
|
for model in models:
|
|
if not hasattr(model, "model"):
|
|
continue
|
|
loaded_model = model_lookup.get(model)
|
|
if loaded_model is None:
|
|
loaded_model = LoadedModel(model)
|
|
model_lookup[model] = loaded_model
|
|
loaded_model.currently_used = True
|
|
loaded.append(loaded_model)
|
|
|
|
# Unload unused models only if necessary
|
|
device = get_torch_device()
|
|
to_remove = []
|
|
if len(current_loaded_models) > 10 or (is_device_cuda(device) and get_cached_memory(device) < 1 * 1024 * 1024 * 1024): # >10 models or <1GB VRAM
|
|
for i, loaded_model in enumerate(current_loaded_models):
|
|
if not loaded_model.currently_used:
|
|
model = loaded_model.model
|
|
if model is None:
|
|
to_remove.append(i)
|
|
continue
|
|
try:
|
|
mem_freed = loaded_model.model_unload()
|
|
to_remove.append(i)
|
|
if hasattr(model, 'detach') and hasattr(model, 'patched_weights') and model.patched_weights:
|
|
model.detach(unpatch_all=True)
|
|
except Exception as e:
|
|
logging.error(f"Failed to unload model at index {i}: {e}")
|
|
for i in reversed(to_remove):
|
|
current_loaded_models.pop(i)
|
|
|
|
lowvram_model_memory = 0
|
|
if vram_state == VRAMState.LOW_VRAM and not force_full_load:
|
|
lowvram_model_memory = max(
|
|
int(get_total_memory(device) * MIN_WEIGHT_MEMORY_RATIO), 400 * 1024 * 1024)
|
|
elif vram_state == VRAMState.NO_VRAM:
|
|
lowvram_model_memory = 1
|
|
|
|
for l in loaded:
|
|
l.currently_used = True
|
|
if l.should_reload_model(force_patch_weights=force_patch_weights) or l.real_model is None:
|
|
mem_needed = l.model_memory_required(device)
|
|
mem_free = get_free_memory(device)
|
|
if DEBUG_ENABLED:
|
|
logging.debug(
|
|
f"Loading {l.model.__class__.__name__}: mem_needed={mem_needed / 1024**3:.2f} GB, free={mem_free / 1024**3:.2f} GB")
|
|
|
|
if mem_free < mem_needed + minimum_memory_required:
|
|
free_memory(mem_needed + minimum_memory_required,
|
|
device, keep_loaded=loaded)
|
|
mem_free = get_free_memory(device)
|
|
|
|
stream = get_offload_stream(device)
|
|
with torch.cuda.stream(stream) if stream is not None else torch.no_grad():
|
|
l.model_load(lowvram_model_memory=lowvram_model_memory, force_patch_weights=force_patch_weights)
|
|
if loaded_model not in current_loaded_models:
|
|
current_loaded_models.append(l) # append for efficiency
|
|
sync_stream(device, stream)
|
|
if DEBUG_ENABLED:
|
|
logging.debug(
|
|
f"Loaded {l.model.__class__.__name__}: free={get_free_memory(device) / 1024**3:.2f} GB")
|
|
|
|
return
|
|
|
|
def load_model_gpu(model):
|
|
"""Load a single model to GPU, wrapper around load_models_gpu."""
|
|
return load_models_gpu([model])
|
|
|
|
|
|
def loaded_models(only_currently_used=False):
|
|
"""Return list of loaded models, optionally only those currently used."""
|
|
output = []
|
|
for m in current_loaded_models:
|
|
if only_currently_used and not m.currently_used:
|
|
continue
|
|
output.append(m.model)
|
|
return output
|
|
|
|
# Data type selection
|
|
def supports_fp8_compute(device=None):
|
|
"""Check if the device supports FP8 computation."""
|
|
if not is_nvidia():
|
|
return False
|
|
if device is None:
|
|
device = get_torch_device()
|
|
props = torch.cuda.get_device_properties(device)
|
|
if props.major >= 9: # Ada Lovelace
|
|
return True
|
|
if props.major == 8 and props.minor >= 9 and torch_version_numeric >= (2, 3):
|
|
if any(platform.win32_ver()) and torch_version_numeric < (2, 4):
|
|
return False
|
|
return True
|
|
return False
|
|
|
|
def supports_dtype(dtype, device):
|
|
"""Check if the device supports the given data type."""
|
|
if dtype == torch.bfloat16:
|
|
if is_nvidia():
|
|
return torch.cuda.get_device_properties(device).major >= 8
|
|
elif is_amd():
|
|
arch = torch.cuda.get_device_properties(device).gcnArchName
|
|
return any(a in arch for a in ["gfx941", "gfx942"])
|
|
return False
|
|
elif dtype in (torch.float16, torch.float32):
|
|
return True
|
|
elif dtype in FLOAT8_TYPES:
|
|
return supports_fp8_compute(device)
|
|
return False
|
|
|
|
def supports_cast(dtype, device):
|
|
"""Check if the device supports casting to the given data type."""
|
|
if dtype == torch.bfloat16:
|
|
return True
|
|
return supports_dtype(dtype, device)
|
|
|
|
def should_use_fp16(device=None, model_params=0, prioritize_performance=True, manual_cast=False):
|
|
"""Determine if FP16 should be used for the device."""
|
|
if device is None:
|
|
device = get_torch_device()
|
|
if FORCE_FP32:
|
|
return False
|
|
if args.force_fp16:
|
|
return supports_cast(torch.float16, device)
|
|
if is_intel_xpu():
|
|
return True
|
|
if is_mlu():
|
|
props = torch.mlu.get_device_properties(device)
|
|
return props.major >= 3
|
|
if is_ascend_npu():
|
|
return False
|
|
if is_amd():
|
|
arch = torch.cuda.get_device_properties(device).gcnArchName
|
|
if any(a in arch for a in ["gfx1030", "gfx1031", "gfx1010", "gfx1011", "gfx1012", "gfx906", "gfx900", "gfx803"]):
|
|
return manual_cast
|
|
return True
|
|
props = torch.cuda.get_device_properties(device)
|
|
if is_nvidia():
|
|
# Prefer FP32 for low VRAM or older GPUs
|
|
total_vram = get_total_memory(device) / (1024**3)
|
|
if total_vram < 5.9 or props.major <= 7: # Turing (7.5) or Pascal (6.x)
|
|
return False
|
|
if any(platform.win32_ver()) and props.major <= 7:
|
|
return manual_cast and torch.cuda.is_bf16_supported()
|
|
if props.major >= 8:
|
|
return True
|
|
return torch.cuda.is_bf16_supported() and manual_cast and (not prioritize_performance or model_params * 4 > get_total_memory(device))
|
|
|
|
def should_use_bf16(device=None, model_params=0, prioritize_performance=True, manual_cast=False):
|
|
"""Determine if BF16 should be used for the device."""
|
|
if device is None:
|
|
device = get_torch_device()
|
|
if args.force_fp16 or FORCE_FP32:
|
|
return False
|
|
if not is_device_cuda(device):
|
|
return False # BF16 not supported on CPU or MPS
|
|
props = torch.cuda.get_device_properties(device)
|
|
if is_nvidia():
|
|
return props.major >= 8 and supports_cast(torch.bfloat16, device)
|
|
elif is_amd():
|
|
arch = props.gcnArchName
|
|
return any(a in arch for a in ["gfx941", "gfx942"]) and supports_cast(torch.bfloat16, device)
|
|
return False
|
|
|
|
def vae_dtype(device=None, model=None):
|
|
"""
|
|
Select appropriate data type for VAE.
|
|
|
|
Args:
|
|
device: PyTorch device (e.g., 'cuda', 'cpu'). Defaults to get_torch_device().
|
|
model: Optional model to check compatibility (not used in this implementation).
|
|
|
|
Returns:
|
|
torch.dtype: Appropriate data type for VAE (e.g., torch.float32, torch.float16, torch.bfloat16).
|
|
"""
|
|
if device is None:
|
|
device = get_torch_device()
|
|
|
|
# Handle CPU case explicitly to avoid CUDA calls
|
|
if device.type == 'cpu':
|
|
logging.debug(f"VAE dtype: torch.float32 (CPU device)")
|
|
return torch.float32
|
|
|
|
# Handle forced FP32/FP16 via command-line arguments
|
|
if args.force_fp32_vae:
|
|
if DEBUG_ENABLED:
|
|
logging.debug(f"VAE dtype: torch.float32 (forced via --force-fp32-vae)")
|
|
return torch.float32
|
|
if args.force_fp16_vae:
|
|
if supports_cast(torch.float16, device):
|
|
if DEBUG_ENABLED:
|
|
logging.debug(f"VAE dtype: torch.float16 (forced via --force-fp16-vae)")
|
|
return torch.float16
|
|
if DEBUG_ENABLED:
|
|
logging.debug(f"VAE dtype: torch.float32 (FP16 not supported on {device})")
|
|
return torch.float32
|
|
|
|
# Handle NVIDIA GPUs
|
|
if is_nvidia():
|
|
props = torch.cuda.get_device_properties(device)
|
|
total_vram = get_total_memory(device) / (1024**3)
|
|
if total_vram < 5.9 or props.major <= 7: # Turing (7.5) or Pascal (6.x)
|
|
# Try FP16 with fallback to FP32 if unstable
|
|
if supports_cast(torch.float16, device) and total_vram >= 3.9:
|
|
if DEBUG_ENABLED:
|
|
logging.debug(f"VAE dtype: torch.float16 (Turing SM {props.major}.{props.minor}, VRAM {total_vram:.1f} GB)")
|
|
return torch.float16
|
|
if DEBUG_ENABLED:
|
|
logging.debug(f"VAE dtype: torch.float32 (Turing SM {props.major}.{props.minor}, low VRAM {total_vram:.1f} GB)")
|
|
return torch.float32
|
|
|
|
# Handle bfloat16 and FP16 for other devices
|
|
if should_use_bf16(device=device, prioritize_performance=False):
|
|
if DEBUG_ENABLED:
|
|
logging.debug(f"VAE dtype: torch.bfloat16 (device supports BF16)")
|
|
return torch.float16
|
|
if should_use_fp16(device=device, prioritize_performance=False):
|
|
if DEBUG_ENABLED:
|
|
logging.debug(f"VAE dtype: torch.float16 (device supports FP16)")
|
|
return torch.float16
|
|
|
|
# Default fallback
|
|
if DEBUG_ENABLED:
|
|
logging.debug(f"VAE dtype: torch.float32 (default fallback)")
|
|
return torch.float32
|
|
|
|
def unet_dtype(device=None, model=None, model_params=None, supported_dtypes=None, weight_dtype=None):
|
|
"""Select appropriate data type for UNet."""
|
|
if device is None:
|
|
device = get_torch_device()
|
|
model_params = module_size(model) // 4 if model is not None else 0
|
|
|
|
# FP8 support
|
|
if args.fp8_e4m3fn_unet and supports_fp8_compute(device):
|
|
return torch.float8_e4m3fn
|
|
if args.fp8_e5m2_unet and supports_fp8_compute(device):
|
|
return torch.float8_e5m2
|
|
fp8_dtype = None
|
|
if weight_dtype in FLOAT8_TYPES:
|
|
fp8_dtype = weight_dtype
|
|
if fp8_dtype is not None:
|
|
if supports_fp8_compute(device):
|
|
return fp8_dtype
|
|
free_model_memory = maximum_vram_for_weights(device)
|
|
if model_params * 2 > free_model_memory:
|
|
return fp8_dtype
|
|
|
|
# Check supported_dtypes and weight_dtype
|
|
if supported_dtypes is not None and weight_dtype is not None:
|
|
for dtype in supported_dtypes:
|
|
if dtype == weight_dtype:
|
|
return dtype
|
|
|
|
# Fallback to bf16/fp16/fp32 based on device and args
|
|
if args.force_fp16 and supports_cast(torch.float16, device):
|
|
return torch.float16
|
|
if args.force_fp32:
|
|
return torch.float32
|
|
if should_use_bf16(device, model_params, prioritize_performance=True):
|
|
return torch.bfloat16
|
|
if should_use_fp16(device, model_params, prioritize_performance=True):
|
|
return torch.float16
|
|
for dt in supported_dtypes:
|
|
if dt == torch.float16 and should_use_fp16(device=device, model_params=model_params, manual_cast=True):
|
|
return torch.float16
|
|
if dt == torch.bfloat16 and should_use_bf16(device, model_params=model_params, manual_cast=True):
|
|
return torch.bfloat16
|
|
return torch.float32
|
|
|
|
def unet_offload_device():
|
|
"""Determine device for UNet offloading (GPU or CPU)."""
|
|
if vram_state == VRAMState.HIGH_VRAM:
|
|
return get_torch_device()
|
|
return torch.device("cpu")
|
|
|
|
def unet_inital_load_device(parameters, dtype):
|
|
"""Determine initial load device for UNet based on model size and dtype."""
|
|
torch_dev = get_torch_device()
|
|
if vram_state in [VRAMState.HIGH_VRAM, VRAMState.SHARED]:
|
|
return torch_dev
|
|
cpu_dev = torch.device("cpu")
|
|
if DISABLE_SMART_MEMORY:
|
|
return cpu_dev
|
|
model_size = dtype_size(dtype) * parameters
|
|
mem_dev = get_free_memory(torch_dev)
|
|
mem_cpu = get_free_memory(cpu_dev)
|
|
if mem_dev > mem_cpu and model_size < mem_dev * 0.8: # 80% threshold
|
|
return torch_dev
|
|
return cpu_dev
|
|
|
|
def unet_manual_cast(weight_dtype, inference_device, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32]):
|
|
"""Determine if manual casting is needed for UNet dtype."""
|
|
# No cast needed for fp32/fp64
|
|
if weight_dtype in [torch.float32, torch.float64]:
|
|
return None
|
|
|
|
# Check FP8 support
|
|
if weight_dtype in [torch.float8_e4m3fn, torch.float8_e5m2] and supports_fp8_compute(inference_device):
|
|
return None
|
|
|
|
# Check FP16 support
|
|
fp16_supported = should_use_fp16(inference_device, prioritize_performance=True)
|
|
if fp16_supported and weight_dtype == torch.float16:
|
|
return None
|
|
|
|
# Check BF16 support
|
|
bf16_supported = should_use_bf16(inference_device)
|
|
if bf16_supported and weight_dtype == torch.bfloat16:
|
|
return None
|
|
|
|
# Prioritize FP16 if supported and in supported_dtypes
|
|
if fp16_supported and torch.float16 in supported_dtypes:
|
|
return torch.float16
|
|
|
|
# Check other supported dtypes
|
|
for dt in supported_dtypes:
|
|
if dt == torch.float16 and fp16_supported:
|
|
return torch.float16
|
|
if dt == torch.bfloat16 and bf16_supported:
|
|
return torch.bfloat16
|
|
if dt in [torch.float8_e4m3fn, torch.float8_e5m2] and supports_fp8_compute(inference_device):
|
|
return dt
|
|
# Fallback to FP32
|
|
return torch.float32
|
|
|
|
|
|
def text_encoder_offload_device():
|
|
"""Determine device for offloading text encoder."""
|
|
return torch.device("cpu") # Keep offload on CPU to save VRAM
|
|
|
|
|
|
def text_encoder_device():
|
|
"""Determine device for text encoder (prefer GPU)."""
|
|
if vram_state in (VRAMState.HIGH_VRAM, VRAMState.NORMAL_VRAM, VRAMState.LOW_VRAM):
|
|
return get_torch_device() # Prefer GPU for low VRAM
|
|
return torch.device("cpu")
|
|
|
|
def text_encoder_initial_device(load_device, offload_device, model_size=0):
|
|
"""Determine initial device for text encoder."""
|
|
if load_device == offload_device or model_size <= 512 * 1024 * 1024:
|
|
return load_device
|
|
if is_device_mps(load_device):
|
|
return load_device
|
|
mem_l = get_free_memory(load_device)
|
|
mem_o = get_free_memory(offload_device)
|
|
if mem_l > (mem_o * 0.5) and model_size * 1.2 < mem_l:
|
|
return load_device
|
|
return offload_device
|
|
|
|
def unet_inital_load_device(parameters, dtype):
|
|
"""Determine initial load device for UNet based on model size and dtype."""
|
|
torch_dev = get_torch_device()
|
|
if vram_state in [VRAMState.HIGH_VRAM, VRAMState.SHARED]:
|
|
return torch_dev
|
|
|
|
cpu_dev = torch.device("cpu")
|
|
if DISABLE_SMART_MEMORY:
|
|
return cpu_dev
|
|
|
|
model_size = dtype_size(dtype) * parameters # Size in bytes
|
|
mem_dev = get_free_memory(torch_dev) # Free VRAM
|
|
mem_cpu = get_free_memory(cpu_dev) # Free RAM
|
|
|
|
# Prefer GPU if VRAM > RAM and model fits in VRAM
|
|
if mem_dev > mem_cpu and model_size < mem_dev * 0.8: # 80% threshold
|
|
return torch_dev
|
|
return cpu_dev
|
|
|
|
def maximum_vram_for_weights(device=None):
|
|
"""Calculate maximum VRAM available for model weights."""
|
|
if device is None:
|
|
device = get_torch_device()
|
|
return (get_total_memory(device) * 0.9 - minimum_inference_memory())
|
|
|
|
|
|
def force_channels_last():
|
|
"""
|
|
Check if channels_last format should be used for tensors.
|
|
Safe for Turing GPUs with FP32 VAE.
|
|
"""
|
|
if args.force_channels_last:
|
|
if DEBUG_ENABLED:
|
|
logging.debug("force_channels_last: Enabled via --force-channels-last")
|
|
return True
|
|
if cpu_state == CPUState.GPU and is_nvidia() and torch.cuda.is_available():
|
|
if DEBUG_ENABLED:
|
|
total_vram = get_total_memory(get_torch_device()) / (1024 * 1024 * 1024) # VRAM in GB
|
|
logging.debug(
|
|
f"force_channels_last: Enabled for NVIDIA GPU with {total_vram:.1f} GB VRAM")
|
|
return True
|
|
logging.debug("force_channels_last: Disabled")
|
|
return False
|
|
|
|
def intermediate_device():
|
|
"""Determine device for intermediate computations (GPU or CPU)."""
|
|
if args.gpu_only:
|
|
return get_torch_device()
|
|
return torch.device("cpu")
|
|
|
|
def get_autocast_device(dev):
|
|
"""Determine device type for autocast (e.g., cuda, cpu, mps)."""
|
|
if hasattr(dev, 'type'):
|
|
return dev.type
|
|
return "cuda"
|
|
|
|
def vae_offload_device():
|
|
"""Determine device for VAE offloading (GPU or CPU)."""
|
|
if args.gpu_only:
|
|
return get_torch_device()
|
|
return torch.device("cpu")
|
|
|
|
def vae_device():
|
|
"""Determine device for VAE (GPU or CPU)."""
|
|
if args.cpu_vae:
|
|
return torch.device("cpu")
|
|
return get_torch_device()
|
|
|
|
def pick_weight_dtype(dtype, fallback_dtype, device=None):
|
|
"""Select appropriate dtype for model weights, using fallback if needed."""
|
|
if dtype is None:
|
|
dtype = fallback_dtype
|
|
elif dtype_size(dtype) > dtype_size(fallback_dtype):
|
|
dtype = fallback_dtype
|
|
if not supports_cast(device, dtype):
|
|
dtype = fallback_dtype
|
|
return dtype
|
|
|
|
def is_device_mps(device):
|
|
"""Check if device is MPS (Apple Silicon)."""
|
|
return isinstance(device, torch.device) and device.type == "mps"
|
|
|
|
def device_supports_non_blocking(device):
|
|
"""Check if device supports non-blocking data transfers."""
|
|
if is_device_mps(device):
|
|
return False # pytorch bug? mps doesn't support non blocking
|
|
if is_intel_xpu():
|
|
return False
|
|
if args.deterministic: # TODO: figure out why deterministic breaks non blocking from gpu to cpu (previews)
|
|
return False
|
|
if directml_enabled:
|
|
return False
|
|
return True
|
|
|
|
def device_should_use_non_blocking(device):
|
|
"""Determine if non-blocking transfers should be used (disabled due to memory issues)."""
|
|
if not device_supports_non_blocking(device):
|
|
return False
|
|
return False
|
|
# return True #TODO: figure out why this causes memory issues on Nvidia and possibly others
|
|
|
|
def text_encoder_dtype(device=None, model=None):
|
|
"""Select appropriate data type for text encoder."""
|
|
if device is None:
|
|
device = get_torch_device()
|
|
model_params = module_size(model) // 4 if model is not None else 0
|
|
# FP8 support (only for CUDA devices)
|
|
if is_device_cuda(device):
|
|
if getattr(args, 'fp8_e4m3fn_text_enc', False) and supports_fp8_compute(device):
|
|
return torch.float8_e4m3fn
|
|
if getattr(args, 'fp8_e5m2_text_enc', False) and supports_fp8_compute(device):
|
|
return torch.float8_e5m2
|
|
# Check model_dtype safely
|
|
model_dtype = getattr(args, 'model_dtype', None)
|
|
if model_dtype is not None:
|
|
dtype_map = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}
|
|
if model_dtype in dtype_map:
|
|
# Only allow BF16 on supported devices
|
|
if model_dtype == "bf16" and not (is_device_cuda(device) and should_use_bf16(device)):
|
|
return torch.float16
|
|
if supports_cast(dtype_map[model_dtype], device):
|
|
return dtype_map[model_dtype]
|
|
# CPU/MPS fallback to FP32
|
|
if not is_device_cuda(device):
|
|
return torch.float32
|
|
# CUDA devices: BF16/FP16 based on device support
|
|
if should_use_bf16(device, model_params, prioritize_performance=True):
|
|
return torch.bfloat16
|
|
if should_use_fp16(device, model_params, prioritize_performance=True):
|
|
return torch.float16
|
|
return torch.float16 # Default to FP16 for GPU
|
|
|
|
|
|
def mac_version():
|
|
"""Get macOS version as a tuple."""
|
|
try:
|
|
return tuple(int(n) for n in platform.mac_ver()[0].split("."))
|
|
except:
|
|
return None
|
|
|
|
# Interrupt handling
|
|
class InterruptProcessingException(Exception):
|
|
pass
|
|
|
|
interrupt_processing_mutex = threading.RLock()
|
|
interrupt_processing = False
|
|
|
|
def interrupt_current_processing(value=True):
|
|
"""Set interrupt flag for processing."""
|
|
global interrupt_processing
|
|
global interrupt_processing_mutex
|
|
with interrupt_processing_mutex:
|
|
interrupt_processing = value
|
|
|
|
def lowvram_enabled():
|
|
"""Check if low VRAM mode is enabled."""
|
|
return vram_state == VRAMState.LOW_VRAM
|
|
|
|
def noram_enabled():
|
|
"""Check if no VRAM mode is enabled."""
|
|
return vram_state == VRAMState.NO_VRAM
|
|
|
|
|
|
def processing_interrupted():
|
|
"""Check if processing is interrupted."""
|
|
global interrupt_processing
|
|
global interrupt_processing_mutex
|
|
with interrupt_processing_mutex:
|
|
return interrupt_processing
|
|
|
|
def throw_exception_if_processing_interrupted():
|
|
"""Throw exception if processing is interrupted."""
|
|
global interrupt_processing
|
|
global interrupt_processing_mutex
|
|
with interrupt_processing_mutex:
|
|
if interrupt_processing:
|
|
interrupt_processing = False
|
|
raise InterruptProcessingException()
|
|
|
|
def controlnet_device():
|
|
"""Determine device for ControlNet (GPU or CPU)."""
|
|
if args.gpu_only:
|
|
return get_torch_device()
|
|
return torch.device("cpu")
|
|
|
|
def controlnet_dtype(device=None, model=None):
|
|
"""Select appropriate data type for ControlNet."""
|
|
if device is None:
|
|
device = get_torch_device()
|
|
model_params = module_size(model) // 4 if model is not None else 0
|
|
if args.force_fp16:
|
|
if supports_cast(torch.float16, device):
|
|
logging.debug(f"ControlNet dtype: torch.float16 (forced via --force-fp16)")
|
|
return torch.float16
|
|
logging.debug(f"ControlNet dtype: torch.float32 (FP16 not supported)")
|
|
return torch.float32
|
|
if args.force_fp32:
|
|
logging.debug(f"ControlNet dtype: torch.float32 (forced via --force-fp32)")
|
|
return torch.float32
|
|
if should_use_bf16(device=device, model_params=model_params, prioritize_performance=False):
|
|
logging.debug(f"ControlNet dtype: torch.bfloat16 (device supports BF16)")
|
|
return torch.bfloat16
|
|
if should_use_fp16(device=device, model_params=model_params, prioritize_performance=False):
|
|
logging.debug(f"ControlNet dtype: torch.float16 (device supports FP16)")
|
|
return torch.float16
|
|
logging.debug(f"ControlNet dtype: torch.float32 (default fallback)")
|
|
return torch.float32
|
|
|
|
def cleanup_models():
|
|
"""Clean up models on finalization."""
|
|
soft_empty_cache(clear=False, caller="cleanup_models")
|
|
# Check memory state after soft clear
|
|
device = get_torch_device()
|
|
cache_key = (device, 'free_memory')
|
|
mem_free_total, _ = _device_cache.get(cache_key, (0, 0))
|
|
if mem_free_total == 0:
|
|
mem_free_total, _ = get_free_memory(device, torch_free_too=True)
|
|
total_vram = get_total_memory(device)
|
|
if mem_free_total < 0.4 * total_vram:
|
|
if PROFILING_ENABLED:
|
|
logging.debug(f"cleanup_models: Insufficient VRAM ({mem_free_total/1024**3:.2f} GB < 20% of {total_vram/1024**3:.2f} GB), forcing aggressive clear")
|
|
soft_empty_cache(clear=True, caller="cleanup_models_aggressive")
|
|
|
|
# Profiling context manager
|
|
@contextlib.contextmanager
|
|
def profile_section(name):
|
|
"""Context manager for profiling code sections."""
|
|
if PROFILING_ENABLED:
|
|
start = time.time()
|
|
if DEBUG_ENABLED:
|
|
stack = [frame for frame in traceback.format_stack(
|
|
limit=10) if "model_management" in frame]
|
|
logging.debug(f"Starting {name}, stack: {''.join(stack)}")
|
|
try:
|
|
yield
|
|
finally:
|
|
logging.debug(f"{name}: {time.time() - start:.3f} s")
|
|
else:
|
|
yield
|
|
|
|
def mac_version():
|
|
"""Get macOS version if running on macOS."""
|
|
if platform.system() == "Darwin":
|
|
try:
|
|
version = platform.mac_ver()[0]
|
|
version_parts = version.split(".")
|
|
return (int(version_parts[0]), int(version_parts[1]))
|
|
except:
|
|
return None
|
|
return None
|
|
|
|
# Additional utilities for memory management
|
|
def get_device_memory_info(device=None):
|
|
"""Get detailed memory information for a device."""
|
|
if device is None:
|
|
device = get_torch_device()
|
|
mem_free_total, mem_free_torch = get_free_memory(device, torch_free_too=True)
|
|
mem_total = get_total_memory(device)
|
|
return {
|
|
"free_total": mem_free_total,
|
|
"free_torch": mem_free_torch,
|
|
"total": mem_total,
|
|
"used": mem_total - mem_free_total
|
|
}
|
|
|
|
def optimize_memory_for_device(device=None):
|
|
"""Optimize memory settings based on device capabilities."""
|
|
if device is None:
|
|
device = get_torch_device()
|
|
total_vram = get_total_memory(device) / (1024 * 1024 * 1024) # VRAM in GB
|
|
global vram_state
|
|
if total_vram < 3.9:
|
|
vram_state = VRAMState.NO_VRAM
|
|
logging.info(f"Low VRAM ({total_vram:.1f} GB), enabling NO_VRAM mode")
|
|
elif total_vram < 7.9:
|
|
vram_state = VRAMState.LOW_VRAM
|
|
logging.info(f"Moderate VRAM ({total_vram:.1f} GB), enabling LOW_VRAM mode")
|
|
else:
|
|
vram_state = VRAMState.NORMAL_VRAM
|
|
logging.info(f"Sufficient VRAM ({total_vram:.1f} GB), using NORMAL_VRAM mode")
|
|
|
|
# Initialize device and memory settings
|
|
try:
|
|
optimize_memory_for_device()
|
|
if PROFILING_ENABLED:
|
|
logging.debug("Memory optimization completed")
|
|
except Exception as e:
|
|
logging.error(f"Failed to optimize memory: {e}")
|
|
vram_state = VRAMState.DISABLED
|
|
|
|
def get_device_cache_state():
|
|
"""Return the current state of _device_cache for logging."""
|
|
return _device_cache |