disable the new NaN check for the CI

3 files changed, 6 insertions, 0 deletions

diff --git a/launch.py b/launch.py
index 715427fd..5afb2956 100644
--- a/launch.py
+++ b/launch.py
@@ -286,6 +286,8 @@ def tests(test_dir):
         sys.argv.append("./test/test_files/empty.pt")

     if "--skip-torch-cuda-test" not in sys.argv:

         sys.argv.append("--skip-torch-cuda-test")

+    if "--disable-nan-check" not in sys.argv:

+        sys.argv.append("--disable-nan-check")

 

     print(f"Launching Web UI in another process for testing with arguments: {' '.join(sys.argv[1:])}")

 

diff --git a/modules/devices.py b/modules/devices.py
index 6f034948..206184fb 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -113,6 +113,9 @@ class NansException(Exception):
 def test_for_nans(x, where):
     from modules import shared
 
+    if shared.cmd_opts.disable_nan_check:
+        return
+
     if not torch.all(torch.isnan(x)).item():
         return
 
diff --git a/modules/shared.py b/modules/shared.py
index 483c4c62..a708f23c 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -64,6 +64,7 @@ parser.add_argument("--sub-quad-chunk-threshold", type=int, help="the percentage
 parser.add_argument("--opt-split-attention-invokeai", action='store_true', help="force-enables InvokeAI's cross-attention layer optimization. By default, it's on when cuda is unavailable.")

 parser.add_argument("--opt-split-attention-v1", action='store_true', help="enable older version of split attention optimization that does not consume all the VRAM it can find")

 parser.add_argument("--disable-opt-split-attention", action='store_true', help="force-disables cross-attention layer optimization")

+parser.add_argument("--disable-nan-check", action='store_true', help="do not check if produced images/latent spaces have nans; useful for running without a checkpoint in CI")

 parser.add_argument("--use-cpu", nargs='+', help="use CPU as torch device for specified modules", default=[], type=str.lower)

 parser.add_argument("--listen", action='store_true', help="launch gradio with 0.0.0.0 as server name, allowing to respond to network requests")

 parser.add_argument("--port", type=int, help="launch gradio with given server port, you need root/admin rights for ports < 1024, defaults to 7860 if available", default=None)