feat(script): add switch node elimination

Author: Chaoses-Ib

script/__init__.py

@@ -125,10 +125,10 @@ def _node_to_assign_st(self, node):
                 # `value is str` doesn't work
                 # TODO: BOOLEAN, not used in any node?
                 if type(value) is str:
-                    args[name] = {'exp': astutil.to_str(value)}
+                    args[name] = {'exp': astutil.to_str(value), 'value': value}
                 else:
                     # int, float
-                    args[name] = {'exp': str(value)}
+                    args[name] = {'exp': str(value), 'value': value}
         if hasattr(v, 'inputs'):
             # If a node's output is not used, it is allowed to have dangling inputs, in which case the link is None.
             # TODO: This breaks the order and arg positions.
@@ -149,7 +149,8 @@ def _node_to_assign_st(self, node):
                     'type': input.type,
                     'move': output_links is None or len(output_links) == 1
                 }
-        args = self._keyword_args_to_positional(v.type, args)
+        args_dict = args
+        args = self._keyword_args_to_positional(v.type, args_dict)
 
         args_of_any_type = [arg for arg in args if arg.get('type', None) == '*']
 
@@ -209,6 +210,7 @@ def _node_to_assign_st(self, node):
 
         c = passes.reroute_elimination(v, args, vars, c)
         c = passes.primitive_node_elimination(v, args, vars, c)
+        c = passes.switch_node_elimination(v, args_dict, args, vars, c)
         return c
 
     def _topological_generations_ordered_dfs(self, end_nodes: Union[list[int], None] = None):

script/passes/__init__.py

@@ -1,7 +1,6 @@
 import re
 
 # TODO: Switch nodes
-# e.g. TomePatchModel, CRLoadLoRA, CLIPSetLastLayer
 # How to prevent var rename?
 # e.g. ModelMergeSimple, CRModelInputSwitch
 
@@ -19,7 +18,26 @@ def primitive_node_elimination(v, args, vars, c):
     assert new_c != c, c
     return new_c
 
+SWITCH_NODES = {
+    'CLIPSetLastLayer': [{'stop_at_clip_layer': -1}],
+    'CR Apply ControlNet': [{'switch': 'Off'}, {'strength': 0}],
+    'CR Load LoRA': [{'switch': 'Off'}, {'strength_model': 0, 'strength_clip': 0}],
+    'TomePatchModel': [{'ratio': 0}],
+}
+def switch_node_elimination(v, args_dict: dict, args, vars, c):
+    switch_inputs = SWITCH_NODES.get(v.type)
+    if switch_inputs is None:
+        return c
+    # print('switch_node_elimination:', v.type, args_dict)
+    for switch_input in switch_inputs:
+        for k, v in switch_input.items():
+            # Only widget values are considered
+            if args_dict[k].get('value') == v:
+                return ''
+    return c
+
 __all__ = [
     'reroute_elimination',
     'primitive_node_elimination',
+    'switch_node_elimination',
 ]