From c87fdcc899246cbbf5cb81777a7b7497b0ba66fc Mon Sep 17 00:00:00 2001
From: strint <xiaoyulink@gmail.com>
Date: Sat, 14 Sep 2024 18:44:54 +0800
Subject: [PATCH] add option for nexfort

---
 .../extras_nodes/nodes_nexfort_booster.py              | 10 +++++++++-
 onediff_comfy_nodes/modules/nexfort/booster_basic.py   |  5 +++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/onediff_comfy_nodes/extras_nodes/nodes_nexfort_booster.py b/onediff_comfy_nodes/extras_nodes/nodes_nexfort_booster.py
index f823e2e3a..2ce01888c 100644
--- a/onediff_comfy_nodes/extras_nodes/nodes_nexfort_booster.py
+++ b/onediff_comfy_nodes/extras_nodes/nodes_nexfort_booster.py
@@ -18,6 +18,12 @@
         "max-optimize:max-autotune:benchmark:low-precision:freezing:cudagraphs": "This is the most aggressive combination of compiler modes. It will deliver the best performance but might slow down the compilation significantly.",
     }
 )
+options = collections.OrderedDict(
+    {
+        '{"inductor.optimize_linear_epilogue": false}': "",
+        '{"triton.fuse_attention_allow_fp16_reduction": false, "inductor.optimize_linear_epilogue": false}': "",
+    }
+)
 
 
 class OneDiffNexfortBooster:
@@ -28,6 +34,7 @@ def INPUT_TYPES(s):
                 "fullgraph": ([False, True],),
                 "dynamic": ([True, False],),
                 "mode": ([mode for mode in compiler_modes.keys()],),
+                "option": ([option for option in options.keys()],),
                 "docs_link": (
                     "STRING",
                     {
@@ -47,11 +54,12 @@ def apply(
         fullgraph=False,
         dynamic=None,
         mode="max-autotune:cudagraphs",
+        option="",
         docs_link=None,
     ):
         return (
             BasicNexFortBoosterExecutor(
-                fullgraph=fullgraph, mode=f"{mode}:cache-all", dynamic=dynamic
+                fullgraph=fullgraph, mode=f"{mode}:cache-all", option=option, dynamic=dynamic
             ),
         )
 
diff --git a/onediff_comfy_nodes/modules/nexfort/booster_basic.py b/onediff_comfy_nodes/modules/nexfort/booster_basic.py
index 3f390f9ab..0d1bfbc75 100644
--- a/onediff_comfy_nodes/modules/nexfort/booster_basic.py
+++ b/onediff_comfy_nodes/modules/nexfort/booster_basic.py
@@ -19,12 +19,17 @@ class BasicNexFortBoosterExecutor(BoosterExecutor):
     def __init__(
         self,
         mode: str = "max-optimize:max-autotune:low-precision",
+        option: str ="",
         fullgraph=False,
         dynamic=True,
     ):
         super().__init__()
+        if isinstance(option, str):
+            import json
+            option = json.loads(option)
         options = {
             "mode": mode,
+            "options": option,
             "dynamic": dynamic,
             "fullgraph": fullgraph,
         }