Add a command line option to disable upcasting in some cross attention ops.

2026-01-23 08:04:27 +08:00 · 2023-01-29 13:12:22 -05:00 · 2023-01-29 13:12:22 -05:00 · 3661e10648
commit 3661e10648
parent 50db297cf6
1 changed files with 4 additions and 0 deletions
--- a/main.py
+++ b/main.py
@ -6,6 +6,10 @@ import threading
 import queue
 import traceback

+if '--dont-upcast-attention' in sys.argv:
+    print("disabling upcasting of attention")
+    os.environ['ATTN_PRECISION'] = "fp16"
+
 import torch

 import nodes