RioShiina/LTX-2.5

1import yaml2import os3import re4import importlib5from copy import deepcopy6from comfy_integration.nodes import NODE_CLASS_MAPPINGS7from chain_injectors import discover_injectors, get_registered_features8from core.settings import FEATURES_CONFIG9 10class WorkflowAssembler:11    def __init__(self, recipe_path, dynamic_values=None):12        self.base_path = os.path.dirname(recipe_path)13        self.dynamic_values = dynamic_values or {}14        self.node_counter = 015        self.workflow = {}16        self.node_map = {}17        18        model_type = self.dynamic_values.get('model_type')19        self._load_injector_config(model_type=model_type)20        21        self.recipe = self._load_and_merge_recipe(os.path.basename(recipe_path), self.dynamic_values)22 23    def _load_injector_config(self, model_type=None):24        self.global_injectors = discover_injectors()25        registered_features = get_registered_features()26 27        order = []28        if model_type and model_type in FEATURES_CONFIG:29            enabled_features = FEATURES_CONFIG[model_type].get('enabled_chains', [])30            for feat in enabled_features:31                if feat in registered_features:32                    chain_key = registered_features[feat]['chain_type']33                else:34                    chain_key = f"dynamic_{feat}_chains"35                if chain_key in self.global_injectors and chain_key not in order:36                    order.append(chain_key)37 38        for chain_key in self.global_injectors.keys():39            if chain_key not in order:40                order.append(chain_key)41 42        self.injector_order = order43 44    def _get_unique_id(self):45        self.node_counter += 146        return str(self.node_counter)47 48    def _get_node_template(self, class_type):49        if class_type not in NODE_CLASS_MAPPINGS:50            raise ValueError(f"Node class '{class_type}' not found. Ensure it's correctly imported in comfy_integration/nodes.py.")51        52        node_class = NODE_CLASS_MAPPINGS[class_type]53        input_types = node_class.INPUT_TYPES()54        55        template = {56            "inputs": {},57            "class_type": class_type,58            "_meta": {"title": node_class.NODE_NAME if hasattr(node_class, 'NODE_NAME') else class_type}59        }60        61        all_inputs = {**input_types.get('required', {}), **input_types.get('optional', {})}62        for name, details in all_inputs.items():63            config = details[1] if len(details) > 1 and isinstance(details[1], dict) else {}64            template["inputs"][name] = config.get("default")65            66        return template67 68    def _load_and_merge_recipe(self, recipe_filename, dynamic_values, search_context_dir=None):69        search_path = search_context_dir or self.base_path70        recipe_path_to_use = os.path.join(search_path, recipe_filename)71 72        if not os.path.exists(recipe_path_to_use):73            raise FileNotFoundError(f"Recipe file not found: {recipe_path_to_use}")74 75        with open(recipe_path_to_use, 'r', encoding='utf-8') as f:76            content = f.read()77 78        for key, value in dynamic_values.items():79            if value is not None:80                content = content.replace(f"{{{{ {key} }}}}", str(value))81        82        main_recipe = yaml.safe_load(content)83        84        merged_recipe = {'nodes': {}, 'connections': [], 'ui_map': {}}85        for key in self.injector_order:86            if key.startswith('dynamic_'):87                merged_recipe[key] = {}88 89        parent_recipe_dir = os.path.dirname(recipe_path_to_use)90        for import_path_template in main_recipe.get('imports', []):91            import_path = import_path_template92            for key, value in dynamic_values.items():93                if value is not None:94                    import_path = import_path.replace(f"{{{{ {key} }}}}", str(value))95            96            try:97                imported_recipe = self._load_and_merge_recipe(import_path, dynamic_values, search_context_dir=parent_recipe_dir)98                merged_recipe['nodes'].update(imported_recipe.get('nodes', {}))99                merged_recipe['connections'].extend(imported_recipe.get('connections', []))100                merged_recipe['ui_map'].update(imported_recipe.get('ui_map', {}))101                for key in self.injector_order:102                    if key in imported_recipe and key.startswith('dynamic_'):103                        merged_recipe[key].update(imported_recipe.get(key, {}))104            except FileNotFoundError:105                print(f"Warning: Optional recipe partial '{import_path}' not found. Skipping.")106 107        merged_recipe['nodes'].update(main_recipe.get('nodes', {}))108        merged_recipe['connections'].extend(main_recipe.get('connections', []))109        merged_recipe['ui_map'].update(main_recipe.get('ui_map', {}))110        for key in self.injector_order:111            if key in main_recipe and key.startswith('dynamic_'):112                merged_recipe[key].update(main_recipe.get(key, {}))113        114        return merged_recipe115 116    def add_node(self, class_type: str, inputs: dict = None, title: str = None) -> str:117        template = self._get_node_template(class_type)118        node_data = deepcopy(template)119        if title:120            node_data['_meta']['title'] = title121        if inputs:122            for k, v in inputs.items():123                node_data['inputs'][k] = v124        node_id = self._get_unique_id()125        self.workflow[node_id] = node_data126        return node_id127 128    def connect(self, from_node: str, from_output_idx: int, to_node: str, to_input_name: str):129        from_id = self.node_map.get(from_node, from_node)130        to_id = self.node_map.get(to_node, to_node)131        if from_id in self.workflow and to_id in self.workflow:132            self.workflow[to_id]['inputs'][to_input_name] = [from_id, int(from_output_idx)]133        else:134            print(f"Warning: Cannot connect '{from_node}' -> '{to_node}'. Node ID not found.")135 136    def assemble(self, ui_values):137        self.ui_values = ui_values138        for name, details in self.recipe['nodes'].items():139            if 'class_type' not in details:140                continue141            class_type = details['class_type']142            match = re.search(r"\{\{\s*(\w+)\s*\}\}", class_type)143            if match:144                placeholder_key = match.group(1)145                if placeholder_key in ui_values and ui_values[placeholder_key] is not None:146                    class_type = ui_values[placeholder_key]147                else:148                    continue149 150            template = self._get_node_template(class_type)151            node_data = deepcopy(template)152            153            unique_id = self._get_unique_id()154            self.node_map[name] = unique_id155 156            if 'title' in details:157                node_data['_meta']['title'] = details['title']158            159            if 'params' in details:160                for param, value in details['params'].items():161                    if '.' in param:162                        parent_param, sub_param = param.split('.', 1)163                        if parent_param in node_data['inputs']:164                            if not isinstance(node_data['inputs'][parent_param], dict):165                                node_data['inputs'][parent_param] = {parent_param: node_data['inputs'][parent_param]} if node_data['inputs'][parent_param] else {}166                            node_data['inputs'][parent_param][sub_param] = value167                    elif param in node_data['inputs']:168                        if isinstance(node_data['inputs'][param], dict) and isinstance(value, dict):169                            node_data['inputs'][param].update(value)170                        else:171                            node_data['inputs'][param] = value172            173            self.workflow[unique_id] = node_data174 175        for ui_key, target in self.recipe.get('ui_map', {}).items():176            if isinstance(target, dict):177                for sub_key, sub_target in target.items():178                    val = None179                    if isinstance(ui_values.get(ui_key), dict):180                        val = ui_values[ui_key].get(sub_key)181                    if val is None:182                        val = ui_values.get(sub_key)183                    if val is not None:184                        sub_target_list = sub_target if isinstance(sub_target, list) else [sub_target]185                        for t in sub_target_list:186                            if isinstance(t, str) and ':' in t:187                                target_name, target_param = t.split(':', 1)188                                if target_name in self.node_map:189                                    self.workflow[self.node_map[target_name]]['inputs'][target_param] = val190            elif ui_key in ui_values and ui_values[ui_key] is not None:191                target_list = target if isinstance(target, list) else [target]192                for t in target_list:193                    if isinstance(t, str) and ':' in t:194                        target_name, target_param = t.split(':', 1)195                        if target_name in self.node_map:196                            self.workflow[self.node_map[target_name]]['inputs'][target_param] = ui_values[ui_key]197 198        for conn in self.recipe.get('connections', []):199            if not isinstance(conn.get('to'), str) or not isinstance(conn.get('from'), str):200                continue201            from_name, from_output_idx = conn['from'].split(':')202            to_name, to_input_name = conn['to'].split(':')203            204            from_id = self.node_map.get(from_name)205            to_id = self.node_map.get(to_name)206            207            if from_id and to_id:208                self.workflow[to_id]['inputs'][to_input_name] = [from_id, int(from_output_idx)]209        210        print("--- [Assembler] Applying dynamic injectors ---")211        recipe_chain_types = {key for key in self.recipe if key.startswith('dynamic_')}212        processing_order = [key for key in self.injector_order if key in recipe_chain_types]213 214        for chain_type in processing_order:215            injector_func = self.global_injectors.get(chain_type)216            if injector_func:217                for chain_key, chain_def in self.recipe.get(chain_type, {}).items():218                    if chain_key in ui_values and ui_values[chain_key]:219                        print(f"  -> Injecting '{chain_type}' for '{chain_key}'...")220                        chain_items = ui_values[chain_key]221                        injector_func(self, chain_def, chain_items)222        223        print("--- [Assembler] Finished applying injectors ---")224        225        return self.workflow