From c361b89026442f3412162657f330d500b803e052 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Tue, 17 Jan 2023 11:04:56 +0300
Subject: [PATCH] disable the new NaN check for the CI

---
 launch.py          | 2 ++
 modules/devices.py | 3 +++
 modules/shared.py  | 1 +
 3 files changed, 6 insertions(+)

diff --git a/launch.py b/launch.py
index 715427fd7..5afb29560 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 6f0349483..206184fbf 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 483c4c620..a708f23c3 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)