add low vram warning

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.

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)

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:]