nan catch as an option

modules/dml/hijack/utils.py

@@ -1,8 +1,11 @@
 import torch
 from typing import Callable
-from installer import log
+from modules.shared import log, opts
 
 def catch_nan(func: Callable[[], torch.Tensor]):
+    if not opts.directml_catch_nan:
+        return func()
+
     tries = 0
     tensor = func()
     while tensor.isnan().sum() != 0 and tries < 10:

modules/shared.py

@@ -290,6 +290,7 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
     "ipex_optimize": OptionInfo(True if devices.backend == "ipex" else False, "Enable IPEX Optimize for Intel GPUs"),
     "ipex_optimize_upscaler": OptionInfo(True if devices.backend == "ipex" else False, "Enable IPEX Optimize for Intel GPUs with Upscalers"),
     "directml_memory_provider": OptionInfo(default_memory_provider, 'DirectML memory stats provider', gr.Radio, {"choices": memory_providers}),
+    "directml_catch_nan": OptionInfo(False, "Retry specific operation when NaN is produced if possible. (makes generation slower)"),
     "openvino_disable_model_caching": OptionInfo(False, "OpenVINO disable model caching"),
     "openvino_hetero_gpu": OptionInfo(False, "OpenVINO use Hetero Device for single inference with multiple devices"),
     "openvino_remove_cpu_from_hetero": OptionInfo(False, "OpenVINO remove CPU from Hetero Device"),