diff --git a/backend/args.py b/backend/args.py index 0302dfc1..6758bb89 100644 --- a/backend/args.py +++ b/backend/args.py @@ -56,6 +56,8 @@ parser.add_argument("--cuda-malloc", action="store_true") parser.add_argument("--cuda-stream", action="store_true") parser.add_argument("--pin-shared-memory", action="store_true") +parser.add_argument("--disable-gpu-warning", action="store_true") + args = parser.parse_known_args()[0] # Some dynamic args that may be changed by webui rather than cmd flags. diff --git a/backend/memory_management.py b/backend/memory_management.py index de56e306..03b4eff2 100644 --- a/backend/memory_management.py +++ b/backend/memory_management.py @@ -929,7 +929,7 @@ def get_free_memory(dev=None, torch_free_too=False): mem_free_torch = mem_free_total else: if directml_enabled: - mem_free_total = 1024 * 1024 * 1024 # TODO + mem_free_total = 1024 * 1024 * 1024 * 2 mem_free_torch = mem_free_total elif is_intel_xpu(): stats = torch.xpu.memory_stats(dev) diff --git a/backend/sampling/sampling_function.py b/backend/sampling/sampling_function.py index a9fbec0c..0c7ffd24 100644 --- a/backend/sampling/sampling_function.py +++ b/backend/sampling/sampling_function.py @@ -10,7 +10,7 @@ import collections from backend import memory_management from backend.sampling.condition import Condition, compile_conditions, compile_weighted_conditions from backend.operations import cleanup_cache -from backend.args import dynamic_args +from backend.args import dynamic_args, args from backend import utils @@ -188,6 +188,20 @@ def calc_cond_uncond_batch(model, cond, uncond, x_in, timestep, model_options): to_batch = to_batch_temp[:1] free_memory = memory_management.get_free_memory(x_in.device) + + if not args.disable_gpu_warning: + free_memory_mb = free_memory / (1024.0 * 1024.0) + safe_memory_mb = 1536.0 + if free_memory_mb < safe_memory_mb: + print(f"\n\n----------------------") + print(f"[Low GPU VRAM Warning] Your current GPU free memory is {free_memory_mb:.2f} MB for this diffusion iteration.") + print(f"[Low GPU VRAM Warning] This number is lower than the safe value of {safe_memory_mb:.2f} MB.") + print(f"[Low GPU VRAM Warning] If you continue the diffusion process, you may cause NVIDIA GPU degradation, and the speed may be extremely slow (about 10x slower).") + print(f"[Low GPU VRAM Warning] To solve the problem, you can set the 'GPU Weight' (on the top of page) to a lower value.") + print(f"[Low GPU VRAM Warning] If you cannot find 'GPU Weight', you can click the 'all' option in the 'UI' area on the left-top corner of the webpage.") + print(f"[Low GPU VRAM Warning] If you want to take the risk of NVIDIA GPU fallback and test the 10x slower speed, you can (but are highly not recommended to) add '--disable-gpu-warning' to CMD flags to remove this warning.") + print(f"----------------------\n\n") + for i in range(1, len(to_batch_temp) + 1): batch_amount = to_batch_temp[:len(to_batch_temp) // i] input_shape = [len(batch_amount) * first_shape[0]] + list(first_shape)[1:]