diff --git a/nodes.py b/nodes.py
index 889ec0f..031f8ee 100644
--- a/nodes.py
+++ b/nodes.py
@@ -149,13 +149,14 @@ class simpleDynamicCFG:
         return {"required": {
                                 "model": ("MODEL",),
                                 "boost" : ("BOOLEAN", {"default": True}),
+                                "color_balance" : ("BOOLEAN", {"default": True}),
                               }}
     RETURN_TYPES = ("MODEL",)
     FUNCTION = "patch"
 
     CATEGORY = "model_patches"
 
-    def patch(self, model, boost):
+    def patch(self, model, boost, color_balance):
         advcfg = advancedDynamicCFG()
-        m = advcfg.patch(model,False,False,"hard" if boost else "soft", boost, 6.86)[0]
-        return (m, )
\ No newline at end of file
+        m = advcfg.patch(model,color_balance,color_balance,"hard" if boost else "soft", boost, 6.86)[0]
+        return (m, )