Add support for FORCE_SOFT_EMPTY_CACHE env var to enable cache clearing for AMD GPUs

Signed-off-by: bigcat88 <bigcat88@icloud.com>
This commit is contained in:
bigcat88 2024-12-19 14:35:29 +02:00
parent 1b3a650f19
commit 51d4abc2cd
No known key found for this signature in database
GPG Key ID: 1F0BF0EC3CF22721

View File

@ -25,6 +25,11 @@ import sys
import platform import platform
import weakref import weakref
import gc import gc
import os
FORCE_SOFT_EMPTY_CACHE = os.getenv("FORCE_SOFT_EMPTY_CACHE", "0") == "1"
class VRAMState(Enum): class VRAMState(Enum):
DISABLED = 0 #No vram present: no need to move models to vram DISABLED = 0 #No vram present: no need to move models to vram
@ -1061,7 +1066,7 @@ def supports_fp8_compute(device=None):
return True return True
def soft_empty_cache(force=False): def soft_empty_cache(force: bool = FORCE_SOFT_EMPTY_CACHE):
global cpu_state global cpu_state
if cpu_state == CPUState.MPS: if cpu_state == CPUState.MPS:
torch.mps.empty_cache() torch.mps.empty_cache()