retry if nan is produced from token merging (directml)

modules/dml/hijack/__init__.py

@@ -5,3 +5,4 @@ import modules.dml.hijack.realesrgan_model
 import modules.dml.hijack.plms
 import modules.dml.hijack.diffusers
 import modules.dml.hijack.transformers
+import modules.dml.hijack.tomesd

modules/dml/hijack/tomesd.py

@@ -0,0 +1,22 @@
+import torch
+import tomesd
+from typing import Type
+from modules.dml.hijack.utils import catch_nan
+
+def make_tome_block(block_class: Type[torch.nn.Module]) -> Type[torch.nn.Module]:
+    class ToMeBlock(block_class):
+        # Save for unpatching later
+        _parent = block_class
+
+        def _forward(self, x: torch.Tensor, context: torch.Tensor = None) -> torch.Tensor:
+            m_a, m_c, m_m, u_a, u_c, u_m = tomesd.patch.compute_merge(x, self._tome_info)
+
+            # This is where the meat of the computation happens
+            x = u_a(self.attn1(m_a(self.norm1(x)), context=context if self.disable_self_attn else None)) + x
+            x = catch_nan(lambda: (u_c(self.attn2(m_c(self.norm2(x)), context=context)) + x))
+            x = u_m(self.ff(m_m(self.norm3(x)))) + x
+
+            return x
+    
+    return ToMeBlock
+tomesd.patch.make_tome_block = make_tome_block

modules/dml/hijack/utils.py

@@ -0,0 +1,15 @@
+import torch
+from typing import Callable
+from installer import log
+
+def catch_nan(func: Callable[[], torch.Tensor]):
+    tries = 0
+    tensor = func()
+    while tensor.isnan().sum() != 0 and tries < 10:
+        if tries == 0:
+            log.warning("NaN is produced. Retry with same values...")
+        tries += 1
+        tensor = func()
+    if tensor.isnan().sum() != 0:
+        log.error("Failed to cover NaN.")
+    return tensor