(data: Union[Dict, List, str], in_nodes: bool, node_type: Optional[str])
| 75 | model_filename_extensions = {'.safetensors', '.ckpt', '.pt', '.pth', '.bin'} |
| 76 | |
| 77 | def recursive_search(data: Union[Dict, List, str], in_nodes: bool, node_type: Optional[str]): |
| 78 | if isinstance(data, dict): |
| 79 | for key, value in data.items(): |
| 80 | type_ = value.get('type') if isinstance(value, dict) else None |
| 81 | recursive_search(value, key == 'nodes' if not in_nodes else in_nodes, type_ if in_nodes and not node_type else node_type) |
| 82 | elif isinstance(data, list): |
| 83 | for item in data: |
| 84 | type_ = item.get('type') if isinstance(item, dict) else None |
| 85 | recursive_search(item, in_nodes, type_ if in_nodes and not node_type else node_type) |
| 86 | elif isinstance(data, str) and '.' in data: |
| 87 | original_filepath = data |
| 88 | normalized_filepath = convert_to_windows_path(original_filepath) if is_windows else convert_to_unix_path(original_filepath) |
| 89 | filename = os.path.basename(data) |
| 90 | |
| 91 | if '.' + original_filepath.split('.')[-1] in model_filename_extensions: |
| 92 | file_names.append(ModelFileWithNodeInfo(filename, original_filepath, normalized_filepath)) |
| 93 | |
| 94 | recursive_search(json_data, False, None) |
| 95 | return file_names |
no test coverage detected